| """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"] |
|
|
