Spaces:
Runtime error
Runtime error
| """ | |
| 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, | |
| ) | |
| 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, | |
| ) | |