diff --git a/bgfilter/alpha_post.py b/bgfilter/alpha_post.py index c0f3634..3b413cf 100644 --- a/bgfilter/alpha_post.py +++ b/bgfilter/alpha_post.py @@ -23,6 +23,7 @@ def suppress_alpha_by_chroma( bg_confidence: np.ndarray, trimap: np.ndarray, settings: AlphaPostSettings, + raw_alpha: np.ndarray | None = None, ) -> np.ndarray: """Pull alpha toward 0 where background-colour confidence is high. @@ -30,6 +31,12 @@ def suppress_alpha_by_chroma( pixels are never thinned. Background colour that survived in hair gaps / hole pockets becomes transparent; this also clears the colour fringe that over-unmixing leaves on those low-alpha pixels, since alpha ~ 0 hides it. + + ``raw_alpha`` is the matting model's own prediction before any constraint. + When given, suppression fades where the model is confident the pixel is + opaque (ramp ``suppress_raw_lo -> suppress_raw_hi``): colour evidence may + only veto the matte where the matte is unsure. Background pockets it rates + low stay suppressible; a background-coloured subject it rates ~1 is kept. """ if not settings.chroma_suppress: return alpha @@ -38,6 +45,10 @@ def suppress_alpha_by_chroma( settings.chroma_suppress_bg_low, settings.chroma_suppress_bg_high, ) + if raw_alpha is not None and settings.suppress_raw_lo < 1.0: + suppress = suppress * ( + 1.0 - _smoothstep(raw_alpha, settings.suppress_raw_lo, settings.suppress_raw_hi) + ) factor = 1.0 - settings.chroma_suppress_strength * suppress out = alpha.astype(np.float32).copy() band = trimap == 128 diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 4611dad..e6aafd9 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -103,10 +103,15 @@ def _run_image( else: seg_mask = None trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap) - alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence) + raw_alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence) - alpha = enforce_trimap(alpha, trimap) - alpha = suppress_alpha_by_chroma(alpha, bg_confidence, trimap, settings.alpha_post) + alpha = enforce_trimap(raw_alpha, trimap) + alpha = suppress_alpha_by_chroma( + alpha, bg_confidence, trimap, settings.alpha_post, + # The matte-confidence gate only makes sense for a real matting prediction; + # a chroma-seeded alpha is itself colour evidence, so no gate there. + raw_alpha=raw_alpha if alpha_source == "vitmatte" else None, + ) alpha = clean_alpha(alpha, trimap, settings.alpha_post) foreground = estimate_foreground_rgb( rgb, alpha, bg_confidence, model, settings.foreground diff --git a/bgfilter/settings.py b/bgfilter/settings.py index f660130..9d2b431 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -54,6 +54,13 @@ class AlphaPostSettings: chroma_suppress_bg_low: float = 0.35 chroma_suppress_bg_high: float = 0.80 chroma_suppress_strength: float = 1.0 + # Matte-confidence gate: fade the suppression where ViTMatte itself is confident + # the pixel is opaque (raw alpha ramps lo -> hi), so colour evidence only vetoes + # the matte where the matte is unsure. Protects background-coloured subjects + # (white shirt on pastel blue, blush on pink) that chroma alone cannot tell from + # background residue. Set suppress_raw_lo >= 1.0 to disable the gate. + suppress_raw_lo: float = 0.85 + suppress_raw_hi: float = 0.98 @dataclass(frozen=True) diff --git a/configs/default.yaml b/configs/default.yaml index bda0325..7da6e81 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -40,6 +40,11 @@ alpha_post: chroma_suppress_bg_low: 0.35 chroma_suppress_bg_high: 0.80 chroma_suppress_strength: 1.0 + # Matte-confidence gate: suppression fades where ViTMatte's raw alpha is high + # (ramp lo -> hi), so colour evidence only vetoes the matte where it is unsure. + # Set suppress_raw_lo: 1.0 to disable. + suppress_raw_lo: 0.85 + suppress_raw_hi: 0.98 foreground: enabled: true