scoring.py

"""
scoring.py — Objective latte art metrics via classical CV.

Five sub-scores, each 0-100:
  contrast    — tonal separation foam vs crema, gated by absolute foam lightness
  flow        — directional structure: gradient orientation dominance + edge
                complexity. A single round blob has low flow even if sharp.
  centering   — pattern centroid vs cup center
  definition  — boundary complexity × edge sharpness. A circle edge is sharp
                but simple; a rosetta edge is sharp AND complex. Both factors
                required for a high score.
  texture     — milk quality, softened: only truly exaggerated bubbles penalise.

Key anti-inflation measures:
  - definition = sharpness × normalised contour complexity (perimeter/area ratio)
    so a smooth blob cannot score high on definition
  - flow uses edge-count density inside foam, not just orientation peaks
  - a single-region foam with low contour complexity is capped at 50 definition
  - presence gate and CURVE=1.35 keep mediocre totals honest
"""

from __future__ import annotations
import cv2
import numpy as np

WEIGHTS = {
    "contrast":   0.25,
    "flow":       0.20,
    "centering":  0.10,
    "definition": 0.30,
    "texture":    0.15,
}

MAX_SIDE = 720
CURVE = 1.35


# ---------------------------------------------------------------- utilities

def _load(path: str) -> np.ndarray:
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"Could not read image: {path}")
    h, w = img.shape[:2]
    scale = MAX_SIDE / max(h, w)
    if scale < 1.0:
        img = cv2.resize(img, (int(w * scale), int(h * scale)),
                         interpolation=cv2.INTER_AREA)
    return img


def _find_cup(img: np.ndarray) -> tuple[int, int, int]:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.medianBlur(gray, 7)
    h, w = gray.shape
    min_r = int(min(h, w) * 0.20)
    max_r = int(min(h, w) * 0.55)
    circles = cv2.HoughCircles(
        gray, cv2.HOUGH_GRADIENT, dp=1.2, minDist=min(h, w),
        param1=120, param2=40, minRadius=min_r, maxRadius=max_r,
    )
    if circles is not None and len(circles[0]) > 0:
        cx, cy, r = circles[0][0]
        return int(cx), int(cy), int(r)
    return w // 2, h // 2, int(min(h, w) * 0.42)


def _crema_mask(img: np.ndarray, cx: int, cy: int, r: int) -> np.ndarray:
    mask = np.zeros(img.shape[:2], dtype=np.uint8)
    cv2.circle(mask, (cx, cy), int(r * 0.86), 255, -1)
    return mask


def _foam_masks(img: np.ndarray,
                surface: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    L = lab[:, :, 0]
    vals = L[surface > 0]
    if vals.size == 0:
        z = np.zeros_like(surface)
        return z, z
    thresh, _ = cv2.threshold(vals, 0, 255,
                               cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    raw = ((L > thresh) & (surface > 0)).astype(np.uint8) * 255
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    clean = cv2.morphologyEx(raw, cv2.MORPH_OPEN, kernel)
    clean = cv2.morphologyEx(clean, cv2.MORPH_CLOSE, kernel)
    return raw, clean


def _clamp(x: float) -> float:
    return float(max(0.0, min(100.0, x)))


def _boundary_complexity(foam: np.ndarray) -> float:
    """Normalised perimeter/area ratio — how complex is the foam boundary?

    A perfect circle has the minimum ratio for its area (isoperimetric).
    Latte art patterns have indentations, lobes, fine lines — much higher ratio.
    Returns a value in [0, 1] where:
      ~0.0  = near-perfect circle (blob, foam dump)
      ~0.5  = heart or simple tulip
      ~1.0  = rosetta, fine rosetta, or multi-region pattern
    """
    contours, _ = cv2.findContours(foam, cv2.RETR_EXTERNAL,
                                   cv2.CHAIN_APPROX_NONE)
    if not contours:
        return 0.0
    total_perim = sum(cv2.arcLength(c, True) for c in contours)
    total_area  = max(foam.sum() / 255.0, 1.0)
    # Circularity-based complexity: for a circle, 4π·area/perim²=1
    # We invert and normalise: more complex = further from circle
    circularity = (4 * np.pi * total_area) / max(total_perim ** 2, 1.0)
    # circularity=1 → blob, circularity→0 → very complex
    # map to [0,1] complexity where 0=blob, 1=complex art
    complexity = _clamp((1.0 - circularity) / 0.92 * 100.0) / 100.0
    # Also reward having multiple regions (e.g. layered tulip lobes)
    region_bonus = min(len(contours) - 1, 4) / 4.0 * 0.3
    return min(1.0, complexity + region_bonus)


# ---------------------------------------------------------------- sub-scores

def _contrast_score(img: np.ndarray, surface: np.ndarray,
                    foam: np.ndarray) -> float:
    L = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)[:, :, 0].astype(np.float32)
    crema = (surface > 0) & (foam == 0)
    fm = (foam > 0)
    if fm.sum() < 200 or crema.sum() < 200:
        return 5.0
    foam_mean = float(L[fm].mean())
    gap = foam_mean - float(L[crema].mean())
    lightness = max(0.0, min(1.0, (foam_mean - 120.0) / 80.0))
    return _clamp((gap / 110.0) * lightness * 100.0)


def _flow_score(img: np.ndarray, foam: np.ndarray,
                surface: np.ndarray, complexity: float) -> float:
    """Directional structure — pattern-type neutral, blob-resistant.

    Combines:
      - orientation dominance (strong preferred direction = intent)
      - edge density inside the foam (fine lines = detail)
      - boundary complexity passed in from _boundary_complexity()
    A round blob has low complexity AND diffuse orientations → low flow.
    A swan has a strong sweep direction AND complex boundary → good flow.
    """
    if foam.sum() == 0:
        return 0.0

    L = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)[:, :, 0].astype(np.float32)
    gx = cv2.Sobel(L, cv2.CV_32F, 1, 0, ksize=3)
    gy = cv2.Sobel(L, cv2.CV_32F, 0, 1, ksize=3)
    mag   = cv2.magnitude(gx, gy)
    angle = cv2.phase(gx, gy, angleInDegrees=True)

    inside = (foam > 0)
    mean_mag = float(mag[inside].mean()) if inside.sum() > 0 else 1.0
    mask = inside & (mag > mean_mag * 0.4)
    if mask.sum() < 50:
        return 10.0

    angles  = angle[mask]
    weights = mag[mask]
    hist, _ = np.histogram(angles, bins=16, range=(0, 360), weights=weights)
    hist = hist / (hist.sum() + 1e-6)

    # Orientation dominance
    peak = float(hist.max())
    mean = float(hist.mean())
    dominance = _clamp((peak / (mean + 1e-6) - 1.0) / 5.0 * 100.0)

    # Edge density inside foam
    edge_density = _clamp(mean_mag / 50.0 * 100.0)

    # Boundary complexity feeds directly in
    complexity_score = complexity * 100.0

    return round(0.35 * dominance + 0.30 * edge_density + 0.35 * complexity_score, 1)


def _centering_score(foam: np.ndarray, cx: int, cy: int, r: int) -> float:
    if foam.sum() == 0:
        return 0.0
    ys, xs = np.nonzero(foam)
    d = np.hypot(xs.mean() - cx, ys.mean() - cy) / max(r, 1)
    return _clamp((1.0 - d / 0.45) * 100.0)


def _definition_score(img: np.ndarray, foam: np.ndarray,
                      complexity: float) -> float:
    """Edge sharpness × boundary complexity.

    A smooth circle can have a sharp edge but its complexity is ~0,
    so definition stays low. Real latte art needs both.
    """
    if foam.sum() == 0:
        return 0.0
    L = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)[:, :, 0].astype(np.float32)
    gx = cv2.Sobel(L, cv2.CV_32F, 1, 0, ksize=3)
    gy = cv2.Sobel(L, cv2.CV_32F, 0, 1, ksize=3)
    grad = cv2.magnitude(gx, gy)
    contours, _ = cv2.findContours(foam, cv2.RETR_EXTERNAL,
                                   cv2.CHAIN_APPROX_NONE)
    boundary = np.zeros_like(foam)
    cv2.drawContours(boundary, contours, -1, 255, 3)
    edge_vals = grad[boundary > 0]
    if edge_vals.size == 0:
        return 0.0
    sharpness = _clamp(float(edge_vals.mean()) / 130.0 * 100.0)
    # Multiply by complexity: blob with sharp edge still scores low
    # Use sqrt so complexity doesn't fully zero out slightly-complex patterns
    complexity_factor = (complexity ** 0.5)
    return round(sharpness * max(0.15, complexity_factor), 1)


def _texture_score(img: np.ndarray, foam_raw: np.ndarray,
                   foam_clean: np.ndarray) -> float:
    if foam_clean.sum() == 0:
        return 65.0
    L = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)[:, :, 0].astype(np.float32)
    interior = cv2.erode(foam_clean,
                         cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15)))
    rough_score = 65.0
    if interior.sum() > 200 * 255:
        lap = cv2.Laplacian(L, cv2.CV_32F, ksize=3)
        rough = float(np.abs(lap[interior > 0]).mean())
        rough_score = _clamp((1.0 - max(0.0, rough - 8.0) / 12.0) * 100.0)

    diff = cv2.bitwise_xor(foam_raw, foam_clean)
    speckle = diff.sum() / max(foam_clean.sum(), 1)
    speckle_score = _clamp((1.0 - max(0.0, speckle - 0.10) / 0.25) * 100.0)
    return round(0.60 * rough_score + 0.40 * speckle_score, 1)


def _presence_factor(foam_frac: float) -> float:
    if foam_frac < 0.02:
        return 0.15
    if foam_frac < 0.08:
        return 0.15 + 0.85 * (foam_frac - 0.02) / 0.06
    if foam_frac <= 0.45:
        return 1.0
    if foam_frac <= 0.65:
        return 1.0 - 0.6 * (foam_frac - 0.45) / 0.20
    return 0.4


# ---------------------------------------------------------------- entry point

def score_image(path: str) -> dict:
    img = _load(path)
    cx, cy, r = _find_cup(img)
    surface = _crema_mask(img, cx, cy, r)
    foam_raw, foam = _foam_masks(img, surface)
    foam_frac = foam.sum() / max(surface.sum(), 1)

    complexity = _boundary_complexity(foam)

    scores = {
        "contrast":   round(_contrast_score(img, surface, foam), 1),
        "flow":       round(_flow_score(img, foam, surface, complexity), 1),
        "centering":  round(_centering_score(foam, cx, cy, r), 1),
        "definition": round(_definition_score(img, foam, complexity), 1),
        "texture":    round(_texture_score(img, foam_raw, foam), 1),
    }
    raw_total = sum(scores[k] * WEIGHTS[k] for k in WEIGHTS)
    gated = raw_total * _presence_factor(float(foam_frac))
    total = 100.0 * (gated / 100.0) ** CURVE

    return {
        "total": round(total, 1),
        "subscores": scores,
        "foam_fraction": round(float(foam_frac), 3),
        "cup": {"cx": cx, "cy": cy, "r": r},
        "weakest": min(scores, key=scores.get),
        "bubbly": scores["texture"] < 40.0 and float(foam_frac) >= 0.04,
        "complexity": round(complexity, 3),
    }


if __name__ == "__main__":
    import json, sys
    print(json.dumps(score_image(sys.argv[1]), indent=2))