""" auto_tagger.py --------------- Heuristic category / style / tag suggestions for an extracted sprite. Why this exists --------------- Every sprite used to arrive as `category="unknown", style="unknown", tags=[]`, leaving the operator to classify each one by hand. At a few hundred sprites that is tedious; at the tens of thousands this project is built for it simply does not happen, and the dataset ships with almost no usable metadata -- which is exactly the metadata a training run needs to filter, balance and condition on. What this is NOT ---------------- This is not a classifier and does not pretend to be one. It is CPU-only arithmetic over shape, palette and provenance -- the same constraint as the rest of the pipeline (no GPU, no model weights, no downloads). It cannot tell a sword from a staff. What it CAN do reliably is the coarse, structural half of the problem: "this came out of a tile grid so it is a tile", "this is a 4-frame strip member so it is an effect frame", "this is a small mostly-green object so it is plausibly foliage". Those are worth a lot when the alternative is `unknown` on every row. Everything it produces is a SUGGESTION: - a suggestion only ever fills a field the user has not set; - below `AUTO_TAG_MIN_CONFIDENCE` the category is left as "unknown" rather than guessed, because a confidently wrong label is worse for a training set than an honest blank; - `auto_tagged` is recorded on the sprite so a reviewer can tell a machine guess from a human decision, and the dataset report can show how much of the metadata was never human-checked. """ from __future__ import annotations from dataclasses import dataclass, field import numpy as np from config import ( AUTO_TAG_LARGE_SPRITE_PX, AUTO_TAG_MIN_CONFIDENCE, AUTO_TAG_TINY_SPRITE_PX, ) @dataclass class TagSuggestion: category: str = "unknown" style: str = "unknown" tags: list[str] = field(default_factory=list) confidence: float = 0.0 reasons: list[str] = field(default_factory=list) def to_dict(self) -> dict: return { "category": self.category, "style": self.style, "tags": self.tags, "confidence": round(self.confidence, 1), "reasons": self.reasons, } # --------------------------------------------------------------------------- # Palette analysis # --------------------------------------------------------------------------- def _opaque_pixels(rgba: np.ndarray) -> np.ndarray: """The sprite's actual artwork pixels, as an (N, 3) RGB array.""" if rgba.size == 0: return np.zeros((0, 3), dtype=np.float32) alpha = rgba[:, :, 3] mask = alpha > 32 if not mask.any(): mask = np.ones(alpha.shape, dtype=bool) return rgba[:, :, :3][mask].astype(np.float32) def _hue_saturation_value(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Vectorised RGB -> HSV (hue in degrees, sat/val in 0..1).""" if rgb.shape[0] == 0: empty = np.zeros((0,), dtype=np.float32) return empty, empty, empty r, g, b = rgb[:, 0] / 255.0, rgb[:, 1] / 255.0, rgb[:, 2] / 255.0 mx = np.maximum(np.maximum(r, g), b) mn = np.minimum(np.minimum(r, g), b) diff = mx - mn hue = np.zeros_like(mx) nz = diff > 1e-6 # Standard piecewise hue, guarded so a fully grey sprite stays at 0. idx = nz & (mx == r) hue[idx] = (60 * ((g[idx] - b[idx]) / diff[idx])) % 360 idx = nz & (mx == g) hue[idx] = (60 * ((b[idx] - r[idx]) / diff[idx]) + 120) % 360 idx = nz & (mx == b) hue[idx] = (60 * ((r[idx] - g[idx]) / diff[idx]) + 240) % 360 sat = np.where(mx > 1e-6, diff / np.maximum(mx, 1e-6), 0.0) return hue, sat, mx def _dominant_hue_family(rgb: np.ndarray) -> tuple[str | None, float]: """ Coarsest possible colour reading: which hue family most of the saturated artwork sits in, and what share of it that family holds. Only reasonably saturated pixels vote -- outlines, shadows and highlights are near-grey and would otherwise drown out the actual colour of a small sprite. """ hue, sat, val = _hue_saturation_value(rgb) if hue.size == 0: return None, 0.0 coloured = (sat > 0.25) & (val > 0.15) if coloured.sum() < max(4, 0.15 * hue.size): return None, 0.0 h = hue[coloured] families = { "red": ((h < 15) | (h >= 345)), "orange": (h >= 15) & (h < 45), "yellow": (h >= 45) & (h < 70), "green": (h >= 70) & (h < 165), "cyan": (h >= 165) & (h < 195), "blue": (h >= 195) & (h < 255), "purple": (h >= 255) & (h < 300), "pink": (h >= 300) & (h < 345), } best_name, best_share = None, 0.0 for name, selector in families.items(): share = float(selector.mean()) if share > best_share: best_name, best_share = name, share return best_name, best_share def _palette_size(rgb: np.ndarray) -> int: if rgb.shape[0] == 0: return 0 return int(np.unique(rgb.astype(np.uint8), axis=0).shape[0]) # --------------------------------------------------------------------------- # Suggestion # --------------------------------------------------------------------------- def _style_from_palette( palette: int, width: int, height: int, rgb: np.ndarray, ) -> tuple[str, str | None]: """ Style, deliberately kept to claims the pixels can actually support. Naming a console is tempting -- "4 colours, must be Game Boy" -- and wrong: a 4-colour dungeon tile is not Game Boy art, and there is no way to recognise an NES palette from one sprite. A style label that confident and that wrong is worse than none, because a training run conditions on it. So only two claims are made: - `gameboy` when the sprite really is a 4-shade MONOCHROME ramp, which is a checkable property, not a guess; - `hd_pixel_art` for a large canvas carrying a wide palette. Everything else stays the honest, generic `pixel_art`. """ longest = max(width, height) if palette <= 4 and rgb.shape[0] > 0: _, sat, _ = _hue_saturation_value(rgb) if sat.size and float(sat.max()) < 0.25: return "gameboy", "4-shade monochrome ramp" if longest >= AUTO_TAG_LARGE_SPRITE_PX and palette > 64: return "hd_pixel_art", "large canvas with a wide palette" return "pixel_art", None def suggest_tags( rgba: np.ndarray, *, segmentation_source: str | None = None, component_count: int = 1, sheet_sprite_count: int = 1, ) -> TagSuggestion: """ Suggest a category, style and descriptive tags for one sprite. `segmentation_source` is the strongest single signal available and comes free from the pipeline: a sprite cut out of a detected tile grid really is a tileset tile, and a sprite cut out of a separator grid on a many-celled sheet really is an animation frame. Neither of those is a guess about content -- it is a fact about how the sprite was produced -- so both are weighted far above anything colour-based. """ height, width = (rgba.shape[0], rgba.shape[1]) if rgba.ndim == 3 else (0, 0) if width == 0 or height == 0: return TagSuggestion() pixels = _opaque_pixels(rgba) palette = _palette_size(pixels) hue_family, hue_share = _dominant_hue_family(pixels) alpha = rgba[:, :, 3] coverage = float((alpha > 32).mean()) longest = max(width, height) shortest = max(1, min(width, height)) aspect = longest / shortest tags: list[str] = [] reasons: list[str] = [] category = "unknown" confidence = 0.0 # --- provenance: the only signals that are facts, not guesses -------- if segmentation_source == "tile": category = "tile" confidence = 85.0 tags.append("tileset") reasons.append("cut from a detected tile grid") if coverage > 0.97: tags.append("terrain") confidence = 90.0 reasons.append("fully opaque, edge to edge") elif segmentation_source == "separator" and sheet_sprite_count >= 4: category = "effect" confidence = 70.0 tags.append("animation_frame") reasons.append(f"one frame of a {sheet_sprite_count}-frame strip") # --- shape ------------------------------------------------------------ if aspect >= 4: tags.append("wide" if width > height else "tall") reasons.append(f"aspect ratio {aspect:.1f}:1") if category == "unknown": category = "decoration" confidence = 45.0 if longest <= AUTO_TAG_TINY_SPRITE_PX: tags.append("tiny") if category == "unknown": category = "item" confidence = 45.0 reasons.append(f"small isolated object ({width}x{height})") elif longest >= AUTO_TAG_LARGE_SPRITE_PX: tags.append("large") if category == "unknown": category = "environment" confidence = 40.0 reasons.append(f"large canvas ({width}x{height})") # --- colour: weakest signal, only ever nudges ------------------------ if hue_family: tags.append(hue_family) if hue_share > 0.6: reasons.append(f"predominantly {hue_family} ({hue_share:.0%})") if category == "unknown" and hue_family == "green": category = "plant" confidence = 40.0 elif category == "unknown" and hue_family in ("cyan", "blue") and longest <= 24: category = "item" confidence = 35.0 if component_count > 3: tags.append("multi_part") style, style_reason = _style_from_palette(palette, width, height, pixels) if style_reason: reasons.append(style_reason) if confidence < AUTO_TAG_MIN_CONFIDENCE: # Below the bar we deliberately hand back "unknown" rather than a # coin flip. The tags still go through -- an unlabelled sprite # tagged "tiny, green" is strictly more useful than a bare one, # and a descriptive tag being wrong costs far less than a wrong # category, which is what a training run filters and balances on. category = "unknown" return TagSuggestion( category=category, style=style, tags=sorted(set(tags)), confidence=confidence, reasons=reasons, )