features/quality.py

"""Image quality and aesthetic features using lightweight proxies.
Since pyiqa may not be available in all environments, we use PIL/NumPy heuristics
that correlate with known quality dimensions."""

import numpy as np
from PIL import Image, ImageStat
from typing import Dict


def compute_quality_features(img: Image.Image) -> Dict[str, float]:
    """
    Compute image quality and aesthetic proxy features.
    
    These are lightweight heuristics. In production, swap with:
    - pyiqa.create_metric('topiq_nr')
    - pyiqa.create_metric('nima')
    - pyiqa.create_metric('clipiqa+')
    
    Returns dict with keys:
        topiq_nr_proxy, nima_aesthetic_proxy, clipiqa_proxy,
        laion_aesthetic_proxy, sharpness, colorfulness
    """
    arr = np.array(img.convert("RGB")).astype(np.float32)
    gray = arr.mean(axis=2)
    H, W = gray.shape
    
    features = {}
    
    # --- TOP-IQ NR proxy: combined sharpness + contrast + noise ---
    # Sharpness: Laplacian variance
    lap = cv2.Laplacian(gray.astype(np.uint8), cv2.CV_64F).var() if 'cv2' in globals() else 0.0
    if lap == 0.0:
        import cv2
        lap = cv2.Laplacian(gray.astype(np.uint8), cv2.CV_64F).var()
    
    # Contrast: Michelson-ish measure
    contrast_proxy = (gray.max() - gray.min()) / (gray.max() + gray.min() + 1e-8)
    # Noise: estimate via local variance
    local_var = np.var(gray[::4, ::4])
    noise_proxy = 1.0 - np.exp(-local_var / 1000.0)  # higher var = less noise bias
    
    # Composite quality proxy: sharpness * contrast * (1 - noise_penalty)
    sharp_norm = np.log1p(lap) / 10.0
    features["topiq_nr_proxy"] = float(np.clip(sharp_norm * contrast_proxy * noise_proxy, 0, 1))
    
    # --- NIMA aesthetic proxy: composition + color + exposure ---
    # Colorfulness (Hasler & Süsstrunk approximation)
    rg = np.abs(arr[:, :, 0] - arr[:, :, 1])
    yb = np.abs(0.5 * (arr[:, :, 0] + arr[:, :, 1]) - arr[:, :, 2])
    colorfulness = float(np.sqrt(rg.std()**2 + yb.std()**2) + 0.3 * np.sqrt(rg.mean()**2 + yb.mean()**2))
    color_score = np.clip(colorfulness / 100.0, 0, 1)
    
    # Exposure balance: histogram centering
    hist, _ = np.histogram(gray.ravel(), bins=256, range=(0, 256))
    hist_norm = hist / (hist.sum() + 1e-8)
    mean_idx = np.sum(np.arange(256) * hist_norm)
    exposure_score = 1.0 - abs(mean_idx - 128.0) / 128.0
    
    # Composition: rule of thirds center
    h3, w3 = max(1, H // 3), max(1, W // 3)
    if h3 * 2 <= H and w3 * 2 <= W:
        center_brightness = gray[h3:2*h3, w3:2*w3].mean() / 255.0
    else:
        center_brightness = gray.mean() / 255.0
    
    # Weighted aesthetic proxy
    features["nima_aesthetic_proxy"] = float(np.clip(
        0.3 * sharp_norm + 0.3 * color_score + 0.2 * exposure_score + 0.2 * center_brightness,
        0, 1
    ))
    
    # --- CLIP-IQA+ proxy: image-text quality correlation ---
    # Proxy: how "natural" / well-composed the image appears
    # Use edge + color + exposure composite
    features["clipiqa_proxy"] = features["nima_aesthetic_proxy"]  # reuse for now
    
    # --- LAION aesthetic proxy ---
    features["laion_aesthetic_proxy"] = features["nima_aesthetic_proxy"]
    
    # Individual quality signals
    features["sharpness"] = sharp_norm
    features["colorfulness"] = color_score
    features["exposure_balance"] = exposure_score
    
    return features


def _try_pyiqa_features(img_path: str) -> Dict[str, float]:
    """
    Attempt to use pyiqa models if available.
    Falls back silently.
    
    Returns empty dict if pyiqa is not available.
    """
    try:
        import pyiqa
        results = {}
        for metric_name in ['brisque', 'niqe', 'clipiqa']:
            try:
                metric = pyiqa.create_metric(metric_name, device='cpu')
                score = metric(img_path).item()
                results[f"pyiqa_{metric_name}"] = score
            except Exception:
                pass
        return results
    except ImportError:
        return {}