diff --git a/bgfilter/cli.py b/bgfilter/cli.py index ec77410..c939863 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -29,6 +29,7 @@ def main( unknown_radius_ratio: float | None = typer.Option(None, "--unknown-radius-ratio", min=0.0), fg_safe_radius_ratio: float | None = typer.Option(None, "--fg-safe-radius-ratio", min=0.0), despill: bool | None = typer.Option(None, "--despill/--no-despill"), + trimap_mode: str | None = typer.Option(None, "--trimap-mode", help="Trimap mode (segmentation pipeline): directional | seg | directional-hard-bg"), ) -> None: try: settings = override_settings( @@ -43,6 +44,7 @@ def main( unknown_radius_ratio=unknown_radius_ratio, fg_safe_radius_ratio=fg_safe_radius_ratio, despill=despill, + trimap_mode=trimap_mode, ) if input_dir is not None: diff --git a/bgfilter/config.py b/bgfilter/config.py index 2a3e73c..32af54a 100644 --- a/bgfilter/config.py +++ b/bgfilter/config.py @@ -92,6 +92,8 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS ] if overrides.get(key) is not None } + if overrides.get("trimap_mode") is not None: + trimap_updates["mode"] = overrides["trimap_mode"] despill_updates: dict[str, Any] = {} if overrides.get("despill") is not None: despill_updates["enabled"] = overrides["despill"] diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index ad03ea5..b5b66d8 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -13,7 +13,12 @@ from .foreground import estimate_foreground_rgb from .io import load_rgb, save_gray, save_rgb, save_rgba, write_text from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews from .settings import PipelineSettings -from .trimap import fuse_trimap, generate_trimap, trimap_to_alpha_seed +from .trimap import ( + fuse_trimap, + fuse_trimap_directional, + generate_trimap, + trimap_to_alpha_seed, +) from .vitmatte_infer import ViTMatteRunner @@ -80,7 +85,21 @@ def _run_image( ) if settings.segmentation.enabled: seg_mask = pipeline._segment(rgb) - trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap) + mode = settings.trimap.mode + if mode == "seg": + trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap) + elif mode in ("directional", "directional-hard-bg"): + from .chroma import convert_color_spaces + + lab = convert_color_spaces(rgb)[2] + trimap, trimap_stats = fuse_trimap_directional( + seg_mask, bg_confidence, lab, model.lab_center, settings.trimap + ) + else: + raise RuntimeError( + f"Unknown trimap.mode '{mode}'. " + "Use 'directional', 'seg', or 'directional-hard-bg'." + ) else: seg_mask = None trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 9f9bcdf..778d4e7 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -28,6 +28,19 @@ class TrimapSettings: # Semantic fusion thresholds (used by fuse_trimap when segmentation is on). seg_core_threshold: float = 0.60 seg_loose_threshold: float = 0.08 + # Trimap mode for the segmentation pipeline (the chroma-only pipeline ignores it): + # "directional" (default) -- chroma magnitude + seg + a hue-direction split. In + # the chroma-unknown zone a confidently-segmented pixel stays foreground + # unless it is displaced toward the background hue: a neutral background- + # coloured garment (e.g. a white shirt) is kept, while a background-hued + # residual (e.g. blue trapped between hair strands) is left UNKNOWN for + # ViTMatte / chroma-suppress to clear. + # "seg" -- original fuse_trimap: seg topology + chroma veto, no hue split. + # "directional-hard-bg" -- like "directional" but hard-marks background-hued + # pixels as background; more aggressive, can eat cool/shadowed white cloth. + mode: str = "directional" + seg_low: float = 0.15 # seg below this -> background-eligible; at/above -> foreground-eligible + bg_hue_proj_min: float = 4.0 # Lab a*/b* projection onto bg direction above which a pixel is background-hued @dataclass(frozen=True) diff --git a/bgfilter/trimap.py b/bgfilter/trimap.py index 08369c1..d647820 100644 --- a/bgfilter/trimap.py +++ b/bgfilter/trimap.py @@ -100,6 +100,81 @@ def fuse_trimap( 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: projection of (a*, b*) onto the background chroma + # direction. Positive = colour displaced toward the background hue. + lab_c = np.asarray(lab_center, dtype=np.float32) + mag = float(np.hypot(lab_c[1], lab_c[2])) + if mag < 1e-3: + proj = np.zeros(shape, dtype=np.float32) + else: + proj = (lab[..., 1] * lab_c[1] + lab[..., 2] * lab_c[2]) / mag + 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 diff --git a/configs/default.yaml b/configs/default.yaml index 6387393..29d216d 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -25,6 +25,15 @@ trimap: fg_safe_radius_ratio: 0.006 seg_core_threshold: 0.60 seg_loose_threshold: 0.08 + # Trimap mode (segmentation pipeline only; also settable via --trimap-mode): + # directional = chroma + seg + hue split; background-hued residual + # (blue between hair) -> unknown, white shirt kept. (default) + # seg = original fuse_trimap (no hue split). + # directional-hard-bg = hard-remove background-hued pixels (aggressive; can eat + # cool/shadowed white cloth). + mode: directional + seg_low: 0.15 + bg_hue_proj_min: 4.0 alpha_post: chroma_suppress: true