Views
No views yet
pip install torch torchvision huggingface_hub opencv-python pillow open-clip-torch1# Define a minimal loader class that matches the uploaded head (512 -> 256 -> 1)
2import torch
3import torch.nn as nn
4from huggingface_hub import PyTorchModelHubMixin
5
6class IQFModel(nn.Module, PyTorchModelHubMixin):
7 def __init__(self, in_dim=512, hidden=256, **kwargs):
8 # Accept either in_dim/hidden or clip_embed_dim/hidden_dim from config.json
9 in_dim = kwargs.pop("clip_embed_dim", in_dim)
10 hidden = kwargs.pop("hidden_dim", hidden)
11 super().__init__()
12 self.mlp = nn.Sequential(
13 nn.Linear(in_dim, hidden),
14 nn.ReLU(),
15 nn.Linear(hidden, 1),
16 )
17 def forward(self, x):
18 return self.mlp(x)
19
20# Load weights from the Hub (defaults to model.safetensors)
21model = IQFModel.from_pretrained("matthewyuan/image-quality-fusion", map_location="cpu")
22model.eval()
23
24# Smoke test on a dummy 512-d vector
25with torch.no_grad():
26 y = model(torch.randn(1, 512)).item()
27print(f"score: {y}")1import torch
2import torch.nn as nn
3from PIL import Image
4import open_clip
5from huggingface_hub import PyTorchModelHubMixin
6
7# Minimal loader class (same as above)
8class IQFModel(nn.Module, PyTorchModelHubMixin):
9 def __init__(self, in_dim=512, hidden=256, **kwargs):
10 in_dim = kwargs.pop("clip_embed_dim", in_dim)
11 hidden = kwargs.pop("hidden_dim", hidden)
12 super().__init__()
13 self.mlp = nn.Sequential(
14 nn.Linear(in_dim, hidden),
15 nn.ReLU(),
16 nn.Linear(hidden, 1),
17 )
18 def forward(self, x):
19 return self.mlp(x)
20
21# 1) Load CLIP ViT-B/32 image encoder (512-d output)
22clip_model, _, clip_preprocess = open_clip.create_model_and_transforms(
23 "ViT-B-32", pretrained="openai"
24)
25clip_model.eval()
26
27# 2) Load the fusion head from the Hub
28fusion = IQFModel.from_pretrained("matthewyuan/image-quality-fusion", map_location="cpu")
29fusion.eval()
30
31def image_to_clip_embedding(img: Image.Image) -> torch.Tensor:
32 x = clip_preprocess(img).unsqueeze(0) # [1, 3, H, W]
33 with torch.no_grad():
34 feat = clip_model.encode_image(x) # [1, 512]
35 feat = feat / feat.norm(dim=-1, keepdim=True)
36 return feat
37
38def predict_quality(image_path: str) -> float:
39 img = Image.open(image_path).convert("RGB")
40 emb = image_to_clip_embedding(img) # [1, 512]
41 with torch.no_grad():
42 score = fusion(emb).item() # scalar
43 return float(score)
44
45print("score:", predict_quality("test.jpg"))| Metric | Value | Description |
|---|---|---|
| Pearson Correlation | 0.520 | Correlation with human judgments |
| R² Score | 0.250 | Coefficient of determination |
| Mean Absolute Error | 1.41 | Average prediction error (1-10 scale) |
| Root Mean Square Error | 1.69 | RMS prediction error |
| Method | Correlation | R² Score | MAE |
|---|---|---|---|
| Fusion Model | 0.520 | 0.250 | 1.41 |
| BRISQUE Only | 0.31 | 0.12 | 2.1 |
| Aesthetic Only | 0.41 | 0.18 | 1.8 |
| CLIP Only | 0.28 | 0.09 | 2.3 |
Input Image (RGB)
├── OpenCV BRISQUE → Technical Quality Score (0-100, normalized)
├── LAION Aesthetic → Aesthetic Score (0-10, normalized)
└── OpenAI CLIP-B32 → Semantic Features (512-dimensional)
↓
Feature Fusion Network
┌─────────────────────────┐
│ BRISQUE: 1D → 64 → 128 │
│ Aesthetic: 1D → 64 → 128│
│ CLIP: 512D → 256 → 128 │
└─────────────────────────┘
↓ (concat)
Deep Fusion Layers (384D → 256D → 128D → 1D)
Dropout (0.3) + ReLU activations
↓
Human-like Quality Score (1.0 - 10.0)1@misc{image-quality-fusion-2024,
2 title={Image Quality Fusion: Multi-Modal Assessment with BRISQUE, Aesthetic, and CLIP Features},
3 author={Matthew Yuan},
4 year={2024},
5 howpublished={\url{https://huggingface.co/matthewyuan/image-quality-fusion}},
6 note={Trained on SPAQ dataset, deployed via GitHub Actions CI/CD}
7}1# Clone repository
2git clone https://github.com/mattkyuan/image-quality-fusion.git
3cd image-quality-fusion
4
5# Install dependencies
6pip install -r requirements.txt
7
8# Run training
9python src/image_quality_fusion/training/train_fusion.py \
10 --image_dir data/images \
11 --annotations data/annotations.csv \
12 --prepare_data \
13 --epochs 50