studio/rigging/segment.py

"""studio.rigging.segment — flood-fill-from-corners cartoon segmentation.

Real alpha channel short-circuits the heuristic: `alpha > 127` is the mask.
Otherwise: connected components of the "background-similar" mask, take the
union of components that touch any of the 4 image corners, invert.

Deterministic, numpy + scipy.ndimage.label, no model imports.
"""
from __future__ import annotations

import numpy as np
from scipy.ndimage import label


DEFAULT_TOLERANCE: int = 15


def flood_fill_segment(
    image: np.ndarray,
    *,
    tolerance: int = DEFAULT_TOLERANCE,
    alpha_threshold: int = 127,
) -> np.ndarray:
    """Return a boolean (H, W) foreground mask.

    image: (H, W) uint8 grayscale, (H, W, 3) RGB uint8, or (H, W, 4) RGBA uint8.
    tolerance: per-channel Chebyshev tolerance for "background-similar" pixels.
    alpha_threshold: cutoff for the RGBA fast path (only used if the alpha
        channel actually varies across the image).
    """
    if image.ndim == 3 and image.shape[-1] == 4:
        alpha = image[..., 3]
        if alpha.min() < 255 and alpha.max() > 0 and alpha.max() != alpha.min():
            return alpha > alpha_threshold
        image = image[..., :3]

    if image.ndim == 2:
        rgb = np.stack([image, image, image], axis=-1)
    elif image.ndim == 3 and image.shape[-1] == 3:
        rgb = image
    else:
        raise ValueError(f"unsupported image shape {image.shape}")

    H, W, _ = rgb.shape
    rgb_i16 = rgb.astype(np.int16)

    corners = np.stack(
        [rgb_i16[0, 0], rgb_i16[0, -1], rgb_i16[-1, 0], rgb_i16[-1, -1]],
        axis=0,
    )
    bg_color = corners.mean(axis=0)
    diff = np.abs(rgb_i16 - bg_color).max(axis=-1)
    similar = diff <= tolerance

    labels, _ = label(similar.astype(np.uint8), structure=np.ones((3, 3), dtype=np.int8))

    bg_label_ids = set()
    for y, x in ((0, 0), (0, W - 1), (H - 1, 0), (H - 1, W - 1)):
        L = int(labels[y, x])
        if L > 0:
            bg_label_ids.add(L)

    if not bg_label_ids:
        return np.ones((H, W), dtype=bool)

    bg_mask = np.isin(labels, np.asarray(sorted(bg_label_ids), dtype=labels.dtype))
    return ~bg_mask


__all__ = ["flood_fill_segment", "DEFAULT_TOLERANCE"]