Spaces:
Runtime error
Runtime error
File size: 2,694 Bytes
d708e6c 2414b69 d708e6c 2414b69 d708e6c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """
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)
|