from __future__ import annotations import numpy as np from .chroma import bg_hue_projection from .deps import require_cv2 from .settings import TrimapSettings def radius_from_ratio(shape: tuple[int, int], ratio: float, minimum: int) -> int: return max(minimum, int(round(max(shape) * ratio))) def elliptical_kernel(radius: int) -> np.ndarray: cv2 = require_cv2() size = radius * 2 + 1 return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) def generate_trimap( bg_confidence: np.ndarray, settings: TrimapSettings ) -> tuple[np.ndarray, dict[str, int]]: cv2 = require_cv2() shape = bg_confidence.shape unknown_radius = radius_from_ratio( shape, settings.unknown_radius_ratio, settings.min_unknown_radius ) fg_safe_radius = radius_from_ratio( shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius ) sure_bg = bg_confidence >= settings.sure_bg_threshold low_bg = bg_confidence <= settings.sure_fg_threshold bg_u8 = sure_bg.astype(np.uint8) unknown_band = cv2.dilate(bg_u8, elliptical_kernel(unknown_radius)).astype(bool) fg_safe = ~cv2.dilate(bg_u8, elliptical_kernel(fg_safe_radius)).astype(bool) sure_fg = low_bg & fg_safe trimap = np.full(shape, 128, dtype=np.uint8) trimap[sure_bg] = 0 trimap[sure_fg] = 255 # Keep a protective unknown band around all sure background, including holes. trimap[unknown_band & ~sure_bg & ~sure_fg] = 128 stats = { "sure_bg_pixels": int((trimap == 0).sum()), "unknown_pixels": int((trimap == 128).sum()), "sure_fg_pixels": int((trimap == 255).sum()), "unknown_radius": int(unknown_radius), "fg_safe_radius": int(fg_safe_radius), } return trimap, stats def fuse_trimap( seg_mask: np.ndarray, bg_confidence: np.ndarray, settings: TrimapSettings ) -> tuple[np.ndarray, dict[str, int]]: """Build a trimap from a semantic subject mask, refined by the chroma key. Authority split: the segmentation mask decides subject *topology* (it keeps colour-contaminated hair as foreground and drops see-through holes), while the chroma key sharpens the flat-background boundary. ViTMatte then refines the unknown band. """ cv2 = require_cv2() shape = seg_mask.shape fg_safe_radius = radius_from_ratio( shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius ) band_radius = max(settings.min_fg_safe_radius, fg_safe_radius // 2) screen_bg = bg_confidence >= settings.sure_bg_threshold # confident flat background core = (seg_mask >= settings.seg_core_threshold).astype(np.uint8) loose = (seg_mask >= settings.seg_loose_threshold).astype(np.uint8) # A small protective band around the mask boundary lets ViTMatte anti-alias # crisp edges; hair gets extra width from BiRefNet's own soft region. The band # is deliberately small so it does not swallow small holes (finger gaps). kernel = elliptical_kernel(band_radius) band = cv2.dilate(loose, kernel).astype(bool) & ~cv2.erode(loose, kernel).astype(bool) # Background follows the mask directly so interior holes stay background, minus # the boundary band; foreground is the confident subject, never the background. sure_fg = core.astype(bool) & ~band & ~screen_bg sure_bg = ((loose == 0) & ~band) | (screen_bg & (loose == 0)) trimap = np.full(shape, 128, dtype=np.uint8) trimap[sure_bg] = 0 trimap[sure_fg] = 255 stats = { "sure_bg_pixels": int((trimap == 0).sum()), "unknown_pixels": int((trimap == 128).sum()), "sure_fg_pixels": int((trimap == 255).sum()), "band_radius": int(band_radius), "fg_safe_radius": int(fg_safe_radius), "seg_subject_pixels": int((core > 0).sum()), } return trimap, stats def fuse_trimap_directional( seg_mask: np.ndarray, bg_confidence: np.ndarray, lab: np.ndarray, lab_center: tuple[float, float, float], settings: TrimapSettings, ) -> tuple[np.ndarray, dict[str, int]]: """Three-signal trimap: chroma magnitude + seg + a directional chroma test. Per pixel (``bgc`` = chroma bg-confidence, ``seg`` = subject confidence):: BG if bgc >= sure_bg_threshold (clean screen colour) OR (seg < seg_low AND bgc > sure_fg_threshold) (seg says not-subject, but never override a pixel chroma is sure is foreground -> keeps fine wisps the segmenter underestimates) FG if bgc <= sure_fg_threshold AND seg >= seg_low (clearly non-screen colour) else (chroma-unknown): if seg >= seg_core_threshold, split by hue -- a pixel *not* displaced toward the background hue (projection < bg_hue_proj_min: e.g. a neutral white shirt) becomes FG, while one strongly displaced toward it (e.g. blue between hair strands) is left unknown for ViTMatte / chroma-suppress. Lower-seg unknowns stay unknown. A thin unknown band at the segmentation silhouette is preserved so ViTMatte can anti-alias the boundary (it is deliberately small so interior holes survive). """ cv2 = require_cv2() shape = seg_mask.shape fg_safe_radius = radius_from_ratio( shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius ) band_radius = max(settings.min_fg_safe_radius, fg_safe_radius // 2) chroma_bg = bg_confidence >= settings.sure_bg_threshold chroma_fg = bg_confidence <= settings.sure_fg_threshold chroma_unknown = ~chroma_bg & ~chroma_fg core = seg_mask >= settings.seg_core_threshold # Directional chroma: positive = colour displaced toward the background hue. proj = bg_hue_projection(lab, lab_center) bg_hued = proj >= settings.bg_hue_proj_min bg = chroma_bg | ((seg_mask < settings.seg_low) & ~chroma_fg) fg = chroma_fg & (seg_mask >= settings.seg_low) rule3_fg = chroma_unknown & core & ~bg_hued if settings.mode == "directional-hard-bg": # Aggressive variant: hard-remove background-hued pixels instead of leaving # them unknown. Can eat cool/shadowed white cloth, so it is not the default. bg = bg | (chroma_unknown & core & bg_hued) trimap = np.full(shape, 128, dtype=np.uint8) trimap[bg] = 0 trimap[fg | rule3_fg] = 255 loose = (seg_mask >= settings.seg_loose_threshold).astype(np.uint8) kernel = elliptical_kernel(band_radius) band = cv2.dilate(loose, kernel).astype(bool) & ~cv2.erode(loose, kernel).astype(bool) trimap[band & ~chroma_bg] = 128 stats = { "sure_bg_pixels": int((trimap == 0).sum()), "unknown_pixels": int((trimap == 128).sum()), "sure_fg_pixels": int((trimap == 255).sum()), "band_radius": int(band_radius), "rule3_fg_pixels": int(rule3_fg.sum()), "bg_hued_pixels": int((bg_hued & chroma_unknown & core).sum()), } return trimap, stats def trimap_to_alpha_seed(trimap: np.ndarray, bg_confidence: np.ndarray) -> np.ndarray: alpha = np.clip(1.0 - bg_confidence, 0.0, 1.0).astype(np.float32) alpha[trimap == 0] = 0.0 alpha[trimap == 255] = 1.0 return alpha