Spaces:
Sleeping
Sleeping
File size: 9,464 Bytes
83b7522 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 |
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from transformers import AutoModel, AutoProcessor
import logging
logger = logging.getLogger(__name__)
class QualityEvaluator:
"""Image quality assessment using multiple SOTA models"""
def __init__(self):
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.models = {}
self.processors = {}
self.load_models()
def load_models(self):
"""Load quality assessment models"""
try:
# Load LAR-IQA model (primary)
logger.info("Loading LAR-IQA model...")
self.load_lar_iqa()
# Load DGIQA model (secondary)
logger.info("Loading DGIQA model...")
self.load_dgiqa()
# Load traditional metrics as fallback
logger.info("Loading traditional quality metrics...")
self.load_traditional_metrics()
except Exception as e:
logger.error(f"Error loading quality models: {str(e)}")
# Use fallback implementation
self.use_fallback_implementation()
def load_lar_iqa(self):
"""Load LAR-IQA model"""
try:
# For now, use a placeholder implementation
# In production, this would load the actual LAR-IQA model
self.models['lar_iqa'] = self.create_mock_model()
self.processors['lar_iqa'] = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
except Exception as e:
logger.warning(f"Could not load LAR-IQA: {str(e)}")
def load_dgiqa(self):
"""Load DGIQA model"""
try:
# Placeholder implementation
self.models['dgiqa'] = self.create_mock_model()
self.processors['dgiqa'] = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
except Exception as e:
logger.warning(f"Could not load DGIQA: {str(e)}")
def load_traditional_metrics(self):
"""Load traditional quality metrics (BRISQUE, NIQE, etc.)"""
try:
# These would be implemented using scikit-image or opencv
self.traditional_metrics_available = True
except Exception as e:
logger.warning(f"Could not load traditional metrics: {str(e)}")
self.traditional_metrics_available = False
def create_mock_model(self):
"""Create a mock model for demonstration purposes"""
class MockQualityModel(nn.Module):
def __init__(self):
super().__init__()
self.backbone = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, 3, padding=1),
torch.nn.ReLU(),
torch.nn.AdaptiveAvgPool2d((1, 1)),
torch.nn.Flatten(),
torch.nn.Linear(64, 1),
torch.nn.Sigmoid()
)
def forward(self, x):
return self.backbone(x) * 10 # Scale to 0-10
model = MockQualityModel().to(self.device)
model.eval()
return model
def use_fallback_implementation(self):
"""Use simple fallback quality assessment"""
logger.info("Using fallback quality assessment implementation")
self.fallback_mode = True
def evaluate_with_lar_iqa(self, image: Image.Image) -> float:
"""Evaluate image quality using LAR-IQA"""
try:
if 'lar_iqa' not in self.models:
return self.fallback_quality_score(image)
# Preprocess image
tensor = self.processors['lar_iqa'](image).unsqueeze(0).to(self.device)
# Get prediction
with torch.no_grad():
score = self.models['lar_iqa'](tensor).item()
return max(0.0, min(10.0, score))
except Exception as e:
logger.error(f"Error in LAR-IQA evaluation: {str(e)}")
return self.fallback_quality_score(image)
def evaluate_with_dgiqa(self, image: Image.Image) -> float:
"""Evaluate image quality using DGIQA"""
try:
if 'dgiqa' not in self.models:
return self.fallback_quality_score(image)
# Preprocess image
tensor = self.processors['dgiqa'](image).unsqueeze(0).to(self.device)
# Get prediction
with torch.no_grad():
score = self.models['dgiqa'](tensor).item()
return max(0.0, min(10.0, score))
except Exception as e:
logger.error(f"Error in DGIQA evaluation: {str(e)}")
return self.fallback_quality_score(image)
def evaluate_traditional_metrics(self, image: Image.Image) -> float:
"""Evaluate using traditional quality metrics"""
try:
# Convert to numpy array
img_array = np.array(image)
# Simple quality metrics based on image statistics
# In production, this would use BRISQUE, NIQE, etc.
# Calculate sharpness (Laplacian variance)
from scipy import ndimage
gray = np.dot(img_array[...,:3], [0.2989, 0.5870, 0.1140])
laplacian_var = ndimage.laplace(gray).var()
sharpness_score = min(10.0, laplacian_var / 100.0)
# Calculate contrast
contrast_score = min(10.0, gray.std() / 25.0)
# Calculate brightness distribution
brightness_score = 10.0 - abs(gray.mean() - 127.5) / 12.75
# Combine scores
quality_score = (sharpness_score * 0.4 +
contrast_score * 0.3 +
brightness_score * 0.3)
return max(0.0, min(10.0, quality_score))
except Exception as e:
logger.error(f"Error in traditional metrics: {str(e)}")
return 5.0 # Default score
def fallback_quality_score(self, image: Image.Image) -> float:
"""Simple fallback quality assessment"""
try:
# Basic quality assessment based on image properties
width, height = image.size
# Resolution score
total_pixels = width * height
resolution_score = min(10.0, total_pixels / 100000.0) # Normalize by 1MP
# Aspect ratio score (prefer standard ratios)
aspect_ratio = width / height
if 0.5 <= aspect_ratio <= 2.0:
aspect_score = 8.0
else:
aspect_score = 5.0
# File format score (prefer lossless)
format_score = 8.0 if image.format == 'PNG' else 6.0
# Combine scores
quality_score = (resolution_score * 0.5 +
aspect_score * 0.3 +
format_score * 0.2)
return max(0.0, min(10.0, quality_score))
except Exception:
return 5.0 # Default neutral score
def evaluate(self, image: Image.Image, anime_mode: bool = False) -> float:
"""
Evaluate image quality using ensemble of models
Args:
image: PIL Image to evaluate
anime_mode: Whether to use anime-specific evaluation
Returns:
Quality score from 0-10
"""
try:
scores = []
# LAR-IQA evaluation
lar_score = self.evaluate_with_lar_iqa(image)
scores.append(lar_score)
# DGIQA evaluation
dgiqa_score = self.evaluate_with_dgiqa(image)
scores.append(dgiqa_score)
# Traditional metrics
traditional_score = self.evaluate_traditional_metrics(image)
scores.append(traditional_score)
# Ensemble scoring
if anime_mode:
# For anime images, weight traditional metrics higher
# as they may be more reliable for stylized content
weights = [0.3, 0.3, 0.4]
else:
# For realistic images, weight modern models higher
weights = [0.4, 0.4, 0.2]
final_score = sum(score * weight for score, weight in zip(scores, weights))
logger.info(f"Quality scores - LAR: {lar_score:.2f}, DGIQA: {dgiqa_score:.2f}, "
f"Traditional: {traditional_score:.2f}, Final: {final_score:.2f}")
return max(0.0, min(10.0, final_score))
except Exception as e:
logger.error(f"Error in quality evaluation: {str(e)}")
return self.fallback_quality_score(image)
|