image_utils.py

"""
image_utils.py
---------------
Shared, low-level image helpers used by every other module.

Hard rule enforced throughout this whole project (see CRITICAL RULES in the
spec): extracted sprite pixels must never be resampled. The only resize
operation allowed anywhere in this codebase is Image.NEAREST, and only for
*preview* purposes -- never for anything that gets written to the dataset.
"""

from __future__ import annotations

import hashlib
import io
from dataclasses import dataclass

import numpy as np
from PIL import Image, ImageFile

# Some scraped sprite sheets are slightly truncated; don't hard-fail on that.
ImageFile.LOAD_TRUNCATED_IMAGES = True


class UnsupportedImageError(Exception):
    """Raised when an uploaded file cannot be safely decoded as an image."""


def load_image_rgba(path_or_bytes) -> Image.Image:
    """
    Load an image and normalize it to RGBA mode WITHOUT resampling pixels.
    Converting palette/RGB/LA modes to RGBA is a lossless, per-pixel exact
    operation (not a resize), so it does not violate pixel preservation.
    """
    try:
        if isinstance(path_or_bytes, (bytes, bytearray)):
            img = Image.open(io.BytesIO(path_or_bytes))
        else:
            img = Image.open(path_or_bytes)
        img.load()
    except Exception as exc:  # noqa: BLE001 - we want to catch any decode failure
        raise UnsupportedImageError(f"Could not decode image: {exc}") from exc

    if img.mode != "RGBA":
        img = img.convert("RGBA")
    return img


def to_numpy(img: Image.Image) -> np.ndarray:
    """RGBA PIL Image -> HxWx4 uint8 numpy array."""
    return np.array(img, dtype=np.uint8)


def to_pil(arr: np.ndarray) -> Image.Image:
    """HxWx4 uint8 numpy array -> RGBA PIL Image."""
    return Image.fromarray(arr.astype(np.uint8), mode="RGBA")


def crop_exact(arr: np.ndarray, x: int, y: int, w: int, h: int) -> np.ndarray:
    """
    Pixel-exact crop, no resampling. Bounds are clamped to the array so a
    slightly-off manual edit never raises.
    """
    height, width = arr.shape[:2]
    x0 = max(0, x)
    y0 = max(0, y)
    x1 = min(width, x + w)
    y1 = min(height, y + h)
    if x1 <= x0 or y1 <= y0:
        return arr[0:0, 0:0]
    return arr[y0:y1, x0:x1].copy()


def pad_nearest(arr: np.ndarray, padding: int) -> np.ndarray:
    """
    Add `padding` pixels of border around a sprite using nearest-pixel
    replication (edge padding), never interpolation. padding=0 returns the
    array unchanged.
    """
    if padding <= 0:
        return arr
    return np.pad(
        arr,
        pad_width=((padding, padding), (padding, padding), (0, 0)),
        mode="edge",
    )


def nearest_preview(img: Image.Image, target_long_side: int = 512) -> Image.Image:
    """
    Upscale a (usually tiny) sprite for on-screen preview using NEAREST only.
    This output must NEVER be written to the dataset -- preview only.
    """
    w, h = img.size
    if w == 0 or h == 0:
        return img
    scale = max(1, int(target_long_side / max(w, h)))
    new_size = (w * scale, h * scale)
    return img.resize(new_size, resample=Image.NEAREST)


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def image_to_png_bytes(img: Image.Image) -> bytes:
    """Serialize losslessly to PNG (PNG is lossless, so this is safe)."""
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=False)
    return buf.getvalue()


@dataclass
class BBox:
    x: int
    y: int
    width: int
    height: int

    @property
    def x2(self) -> int:
        return self.x + self.width

    @property
    def y2(self) -> int:
        return self.y + self.height

    @property
    def area(self) -> int:
        return self.width * self.height

    @property
    def cx(self) -> float:
        return self.x + self.width / 2.0

    @property
    def cy(self) -> float:
        return self.y + self.height / 2.0

    def gap_to(self, other: "BBox") -> tuple[float, float]:
        """
        Signed horizontal/vertical gap between two bounding boxes.
        Negative values mean the boxes overlap on that axis.
        """
        dx = max(other.x - self.x2, self.x - other.x2)
        dy = max(other.y - self.y2, self.y - other.y2)
        return dx, dy

    def to_dict(self) -> dict:
        return {"x": self.x, "y": self.y, "width": self.width, "height": self.height}

    def union(self, other: "BBox") -> "BBox":
        x0 = min(self.x, other.x)
        y0 = min(self.y, other.y)
        x1 = max(self.x2, other.x2)
        y1 = max(self.y2, other.y2)
        return BBox(x0, y0, x1 - x0, y1 - y0)