background.py

"""
background.py
--------------
Border-connected background detection.

The critical rule from the spec: we must NOT treat "detected the background
color" as "make every pixel of that color transparent". A sprite can
legitimately contain the same color as the background inside its own body
(e.g. a black outline on a black-background sheet). Only background pixels
that are *reachable from the sheet's border* through same-colored
(tolerance-aware) neighbours are background. A black pixel enclosed inside
a sprite's silhouette is never touched.

Algorithm
---------
1. Sample border pixels to estimate the background color (mode of the
   border ring, alpha-aware).
2. Flood-fill (BFS, 4-connected) starting from every border pixel that
   matches the background color within tolerance, expanding only through
   neighbours that also match. This produces a boolean "is background"
   mask that is by construction connected to the border.
3. Anything not reached by the flood fill is foreground, even if it is the
   exact same color as the background.
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass

import numpy as np

from config import BACKGROUND_COLOR_TOLERANCE


@dataclass
class BackgroundInfo:
    color: tuple[int, int, int, int]
    is_transparent: bool  # sheet already had a real alpha channel background
    tolerance: int
    mask: np.ndarray  # HxW bool, True = background


def _estimate_border_color(rgb: np.ndarray, alpha: np.ndarray) -> tuple[int, int, int, int]:
    """Mode color of the outermost ring of pixels, alpha-aware."""
    h, w = alpha.shape
    ring_mask = np.zeros((h, w), dtype=bool)
    ring_mask[0, :] = True
    ring_mask[-1, :] = True
    ring_mask[:, 0] = True
    ring_mask[:, -1] = True

    ring_rgb = rgb[ring_mask]
    ring_alpha = alpha[ring_mask]

    # If the border is overwhelmingly transparent, background = transparent.
    if (ring_alpha < 8).mean() > 0.6:
        return (0, 0, 0, 0)

    # Otherwise take the most common RGB triple on the border.
    colors, counts = np.unique(ring_rgb.reshape(-1, 3), axis=0, return_counts=True)
    mode_color = colors[np.argmax(counts)]
    return (int(mode_color[0]), int(mode_color[1]), int(mode_color[2]), 255)


def detect_background(rgba: np.ndarray, tolerance: int = BACKGROUND_COLOR_TOLERANCE) -> BackgroundInfo:
    """
    Detect the background of a sprite sheet using the border-connected
    flood-fill approach described above.
    """
    h, w = rgba.shape[:2]
    rgb = rgba[:, :, :3].astype(np.int16)
    alpha = rgba[:, :, 3]

    bg_color = _estimate_border_color(rgb, alpha)
    is_transparent_bg = bg_color[3] == 0

    if is_transparent_bg:
        # Background = fully (or near-fully) transparent pixels.
        candidate = alpha < 8
    else:
        bg_rgb = np.array(bg_color[:3], dtype=np.int16)
        # Per-pixel color distance (max-channel-diff is cheap and works well
        # for the flat, dithering-free colors typical of sprite sheet mattes).
        diff = np.abs(rgb - bg_rgb).max(axis=2)
        candidate = (diff <= tolerance) & (alpha > 0)
        # Also treat fully transparent pixels as background even on an
        # otherwise opaque-background sheet (mixed sheets do happen).
        candidate |= alpha < 8

    # Border-connected flood fill (BFS) restricted to `candidate` pixels.
    visited = np.zeros((h, w), dtype=bool)
    dq: deque[tuple[int, int]] = deque()

    border_coords = []
    border_coords.extend((0, x) for x in range(w))
    border_coords.extend((h - 1, x) for x in range(w))
    border_coords.extend((y, 0) for y in range(h))
    border_coords.extend((y, w - 1) for y in range(h))

    for y, x in border_coords:
        if candidate[y, x] and not visited[y, x]:
            visited[y, x] = True
            dq.append((y, x))

    while dq:
        cy, cx = dq.popleft()
        for ny, nx in ((cy - 1, cx), (cy + 1, cx), (cy, cx - 1), (cy, cx + 1)):
            if 0 <= ny < h and 0 <= nx < w and not visited[ny, nx] and candidate[ny, nx]:
                visited[ny, nx] = True
                dq.append((ny, nx))

    return BackgroundInfo(
        color=bg_color,
        is_transparent=is_transparent_bg,
        tolerance=tolerance,
        mask=visited,
    )


def foreground_mask(rgba: np.ndarray, bg_info: BackgroundInfo) -> np.ndarray:
    """Convenience: inverse of the background mask."""
    return ~bg_info.mask


def apply_transparency(rgba: np.ndarray, bg_info: BackgroundInfo) -> np.ndarray:
    """
    Return a copy of `rgba` with ONLY the border-connected background
    pixels made transparent. Pixels inside sprites that merely share the
    background color are left completely untouched (both RGB and alpha).
    """
    out = rgba.copy()
    out[bg_info.mask, 3] = 0
    return out