sprite_extractor.py

"""
sprite_extractor.py
--------------------
Turns segmentation results (bounding boxes) into actual exported sprite
images. This is where "extract each sprite at the smallest possible
bounding box, no resize, optional NEAREST padding" is implemented.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from background import BackgroundInfo, apply_transparency
from image_utils import BBox, crop_exact, pad_nearest, to_pil
from segmentation import DetectedSprite


@dataclass
class ExtractedSprite:
    sprite_index: int
    bbox: BBox
    rgba: np.ndarray          # exact pixels, background untouched
    rgba_transparent_bg: np.ndarray  # same pixels, background made transparent
    confidence: float
    confidence_reasons: list[str]
    component_count: int
    source_file: str
    source: str = "gap"   # which segmentation path produced it


def extract_sprite(
    sheet_rgba: np.ndarray,
    bg_info: BackgroundInfo,
    detected: DetectedSprite,
    source_file: str,
    padding: int = 0,
) -> ExtractedSprite:
    """
    Crop exactly `detected.bbox` out of the sheet (no resampling). Produces
    two variants: background kept as-is, and background made transparent
    (only the border-connected background pixels, per background.py).
    """
    b = detected.bbox
    kept = crop_exact(sheet_rgba, b.x, b.y, b.width, b.height)

    # Build the transparent-background version by cropping from a
    # full-sheet transparent copy, so the exact same border-connected mask
    # logic applies consistently regardless of crop boundaries.
    full_transparent = apply_transparency(sheet_rgba, bg_info)
    transparent = crop_exact(full_transparent, b.x, b.y, b.width, b.height)

    if padding > 0:
        kept = pad_nearest(kept, padding)
        transparent = pad_nearest(transparent, padding)

    return ExtractedSprite(
        sprite_index=detected.sprite_index,
        bbox=b,
        rgba=kept,
        rgba_transparent_bg=transparent,
        confidence=detected.confidence,
        confidence_reasons=detected.confidence_reasons,
        component_count=detected.component_count,
        source_file=source_file,
        source=detected.source,
    )


def extract_all(
    sheet_rgba: np.ndarray,
    bg_info: BackgroundInfo,
    detections: list[DetectedSprite],
    source_file: str,
    padding: int = 0,
) -> list[ExtractedSprite]:
    return [
        extract_sprite(sheet_rgba, bg_info, d, source_file, padding=padding)
        for d in detections
    ]


def sprite_to_pil(sprite: ExtractedSprite, keep_background: bool):
    arr = sprite.rgba if keep_background else sprite.rgba_transparent_bg
    return to_pil(arr)