scoring.py

"""
scoring.py
----------
Per-sprite quality analysis, run AFTER segmentation/extraction.

All metrics are CPU-only (numpy / scipy / basic stats), no ML models, per
the hardware constraints (8 vCPU / 32GB RAM, GPU not guaranteed).

Every sub-score is 0-100. The overall score is a weighted combination.
Small sprites are explicitly NOT penalized for being small -- resolution
is only one factor among several, as required by the spec.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy import ndimage


@dataclass
class QualityReport:
    overall: float
    pixel_art_probability: float
    sharpness: float
    color_consistency: float
    compression_artifact_score: float
    transparency_quality: float
    resolution_suitability: float
    pixel_grid_consistency: float
    noise_score: float
    unique_colors: int
    has_alpha: bool

    def to_dict(self) -> dict:
        return {
            "overall_quality_score": round(self.overall, 1),
            "pixel_art_probability": round(self.pixel_art_probability, 1),
            "sharpness_score": round(self.sharpness, 1),
            "color_consistency": round(self.color_consistency, 1),
            "compression_artifact_score": round(self.compression_artifact_score, 1),
            "transparency_quality": round(self.transparency_quality, 1),
            "resolution_suitability": round(self.resolution_suitability, 1),
            "pixel_grid_consistency": round(self.pixel_grid_consistency, 1),
            "noise_score": round(self.noise_score, 1),
            "unique_colors": self.unique_colors,
            "has_alpha": self.has_alpha,
        }


def _unique_colors(rgba: np.ndarray) -> int:
    if rgba.size == 0:
        return 0
    flat = rgba.reshape(-1, 4)
    return int(np.unique(flat, axis=0).shape[0])


def _sharpness_score(rgb: np.ndarray) -> float:
    """
    Laplacian-variance based sharpness. Pixel art is inherently sharp
    (hard edges); a low score suggests blur (e.g. from a bad scrape/resize
    upstream, or JPEG blur).
    """
    if rgb.shape[0] < 3 or rgb.shape[1] < 3:
        return 100.0  # too small to meaningfully measure; don't penalize
    gray = rgb.mean(axis=2)
    lap = ndimage.laplace(gray)
    variance = float(lap.var())
    # Empirically, pixel art laplacian variance is high (hard edges).
    # Normalize with a soft cap.
    score = min(100.0, (variance / 40.0) * 100.0)
    return max(0.0, score)


def _color_consistency_score(rgb: np.ndarray, alpha: np.ndarray) -> float:
    """
    Pixel art typically uses a small, flat palette (little to no gradient
    noise within a single color region). We approximate "flatness" by
    checking what fraction of colors are exact-duplicates of a small
    palette vs. a huge spread of near-unique colors (a sign of photographic
    content, JPEG ringing, or gradients).
    """
    mask = alpha > 10
    if mask.sum() == 0:
        return 50.0
    pixels = rgb[mask]
    total = pixels.shape[0]
    unique = np.unique(pixels, axis=0).shape[0]
    ratio = unique / max(1, total)
    # Lower ratio (few unique colors relative to pixel count) = more
    # consistent/flat = higher score.
    score = 100.0 * (1.0 - min(1.0, ratio * 3.0))
    return max(0.0, min(100.0, score))


def _compression_artifact_score(rgb: np.ndarray, alpha: np.ndarray) -> float:
    """
    Detects JPEG-style 8x8 blockiness / ringing near edges -- pixel art
    should have none of this if it was saved as PNG. We look for small
    high-frequency noise adjacent to hard edges, which flat pixel art does
    not exhibit but re-compressed JPEG pixel art does.
    """
    if rgb.shape[0] < 8 or rgb.shape[1] < 8:
        return 100.0
    gray = rgb.mean(axis=2)
    gx = np.abs(np.diff(gray, axis=1))
    gy = np.abs(np.diff(gray, axis=0))
    # Ringing shows up as many small (1-10 intensity) differences rather
    # than a few large, clean edge jumps typical of flat-color pixel art.
    small_diffs = ((gx > 1) & (gx < 10)).mean() + ((gy > 1) & (gy < 10)).mean()
    score = 100.0 - min(100.0, small_diffs * 150.0)
    return max(0.0, score)


def _transparency_quality_score(alpha: np.ndarray) -> float:
    """
    Good pixel art alpha is binary-ish (fully opaque or fully transparent)
    or has small, deliberate anti-aliased edges. Large regions of
    "mushy" partial alpha (e.g. 40-180 out of 255) suggest a bad cutout or
    a blurred/resized source.
    """
    if alpha.size == 0:
        return 0.0
    if (alpha == 255).all():
        return 100.0  # fully opaque sprite, nothing to penalize
    mushy = ((alpha > 20) & (alpha < 235)).mean()
    score = 100.0 - min(100.0, mushy * 200.0)
    return max(0.0, score)


def _resolution_suitability_score(width: int, height: int) -> float:
    """
    Resolution is only ONE factor and small sprites are explicitly valid
    (8x8, 16x16 etc. per spec). We only penalize truly degenerate sizes
    (near-zero) or absurdly large "sprites" that are probably a
    segmentation failure.
    """
    longest = max(width, height)
    if longest < 2:
        return 0.0
    if longest > 1024:
        return 40.0  # plausible tileset/background slipped through, flag lightly
    return 100.0


def _pixel_grid_consistency_score(rgb: np.ndarray, alpha: np.ndarray) -> float:
    """
    Heuristic for "is this drawn on a clean pixel grid" by checking if
    runs of identical adjacent pixels are chunky (pixel-art-like) rather
    than constantly varying by 1-2 intensity levels (photographic/
    gradient-like). We measure the average run-length of identical pixels
    along rows.
    """
    mask = alpha > 10
    if mask.sum() < 4:
        return 50.0
    gray = (rgb.mean(axis=2)).astype(np.int16)
    row_run_lengths = []
    for row in gray:
        if row.size < 2:
            continue
        changes = np.count_nonzero(np.diff(row))
        run_estimate = row.size / max(1, changes + 1)
        row_run_lengths.append(run_estimate)
    if not row_run_lengths:
        return 50.0
    avg_run = float(np.mean(row_run_lengths))
    # Longer average runs = flatter/blockier = more pixel-art-like.
    score = min(100.0, (avg_run / 3.0) * 100.0)
    return max(0.0, score)


def _noise_score(rgb: np.ndarray, alpha: np.ndarray) -> float:
    """
    High score = low noise. Pixel art has large flat-color regions with
    essentially zero internal variance; noise/dithering/JPEG artifacts show
    up as intensity jitter *within* those flat regions. Raw local variance
    over the whole sprite is dominated by legitimate hard edges (a red
    block next to a white block is a huge, correct jump, not noise), so we
    first find "flat interior" pixels -- ones whose immediate neighbours
    are all nearly identical to them -- and only measure variance among
    those. A photographic/gradient sprite has almost no such flat interior
    and will score low; clean pixel art is mostly flat interior and scores
    high.
    """
    if rgb.shape[0] < 3 or rgb.shape[1] < 3:
        return 100.0
    gray = rgb.mean(axis=2)
    local_max = ndimage.maximum_filter(gray, size=3)
    local_min = ndimage.minimum_filter(gray, size=3)
    local_range = local_max - local_min

    mask = alpha > 10
    if mask.sum() == 0:
        return 50.0

    # "Flat interior" = pixels where the 3x3 neighbourhood barely varies at
    # all (well inside a single flat-colored region, away from any edge).
    flat_interior = mask & (local_range < 6)
    if flat_interior.sum() < max(4, int(0.05 * mask.sum())):
        # Almost no flat regions at all -- e.g. a 1-2px sprite, or genuinely
        # noisy/gradient content with no flat areas. Fall back to a mild
        # penalty based on overall local range rather than claiming 100.
        avg_range = float(local_range[mask].mean()) if mask.sum() else 0.0
        return max(0.0, 100.0 - min(100.0, avg_range * 2.0))

    # Within the flat interior, real dithering/JPEG noise still causes
    # small (1-4 intensity level) jitter even though the *range* filter
    # above didn't catch it at the 6-level threshold. Measure that.
    local_mean = ndimage.uniform_filter(gray, size=3)
    local_sq_mean = ndimage.uniform_filter(gray ** 2, size=3)
    local_var = np.clip(local_sq_mean - local_mean ** 2, 0, None)
    avg_var = float(local_var[flat_interior].mean())
    score = 100.0 - min(100.0, avg_var * 4.0)
    return max(0.0, score)


def _pixel_art_probability(
    sharpness: float,
    color_consistency: float,
    compression: float,
    grid_consistency: float,
    noise: float,
) -> float:
    weights = {
        "sharpness": 0.25,
        "color_consistency": 0.25,
        "compression": 0.2,
        "grid_consistency": 0.2,
        "noise": 0.1,
    }
    return (
        sharpness * weights["sharpness"]
        + color_consistency * weights["color_consistency"]
        + compression * weights["compression"]
        + grid_consistency * weights["grid_consistency"]
        + noise * weights["noise"]
    )


def analyze_sprite(rgba: np.ndarray) -> QualityReport:
    if rgba.size == 0 or rgba.shape[0] == 0 or rgba.shape[1] == 0:
        return QualityReport(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, False)

    rgb = rgba[:, :, :3].astype(np.float32)
    alpha = rgba[:, :, 3]
    has_alpha = bool((alpha < 250).any())

    sharpness = _sharpness_score(rgb)
    color_consistency = _color_consistency_score(rgb, alpha)
    compression = _compression_artifact_score(rgb, alpha)
    transparency = _transparency_quality_score(alpha)
    resolution = _resolution_suitability_score(rgba.shape[1], rgba.shape[0])
    grid_consistency = _pixel_grid_consistency_score(rgb, alpha)
    noise = _noise_score(rgb, alpha)
    unique_colors = _unique_colors(rgba)

    pixel_art_prob = _pixel_art_probability(sharpness, color_consistency, compression, grid_consistency, noise)

    overall = (
        pixel_art_prob * 0.35
        + sharpness * 0.15
        + color_consistency * 0.15
        + compression * 0.1
        + transparency * 0.1
        + resolution * 0.1
        + noise * 0.05
    )

    return QualityReport(
        overall=overall,
        pixel_art_probability=pixel_art_prob,
        sharpness=sharpness,
        color_consistency=color_consistency,
        compression_artifact_score=compression,
        transparency_quality=transparency,
        resolution_suitability=resolution,
        pixel_grid_consistency=grid_consistency,
        noise_score=noise,
        unique_colors=unique_colors,
        has_alpha=has_alpha,
    )