File size: 2,223 Bytes
7faaef2 | 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 | """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"]
|