diff --git a/bgfilter/alpha_post.py b/bgfilter/alpha_post.py index 9e1f0b6..5151c32 100644 --- a/bgfilter/alpha_post.py +++ b/bgfilter/alpha_post.py @@ -13,6 +13,38 @@ def enforce_trimap(alpha: np.ndarray, trimap: np.ndarray) -> np.ndarray: return out +def _smoothstep(x: np.ndarray, edge0: float, edge1: float) -> np.ndarray: + t = np.clip((x - edge0) / max(edge1 - edge0, 1e-6), 0.0, 1.0) + return t * t * (3.0 - 2.0 * t) + + +def suppress_alpha_by_chroma( + alpha: np.ndarray, + bg_confidence: np.ndarray, + trimap: np.ndarray, + settings: AlphaPostSettings, +) -> np.ndarray: + """Pull alpha toward 0 where background (green) confidence is high. + + Only the unknown band (trimap == 128) is touched, so confidently-foreground + pixels are never thinned. Green that survived in hair gaps / hole pockets + becomes transparent; this also clears the magenta fringe that over-unmixing + leaves on those low-alpha pixels, since alpha ~ 0 hides the foreground colour. + """ + if not settings.chroma_suppress: + return alpha + suppress = _smoothstep( + np.clip(bg_confidence, 0.0, 1.0), + settings.chroma_suppress_bg_low, + settings.chroma_suppress_bg_high, + ) + factor = 1.0 - settings.chroma_suppress_strength * suppress + out = alpha.astype(np.float32).copy() + band = trimap == 128 + out[band] = out[band] * factor[band] + return np.clip(out, 0.0, 1.0) + + def _remove_small_components(mask: np.ndarray, min_area: int) -> np.ndarray: cv2 = require_cv2() count, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8) diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 005903a..298d50f 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -6,7 +6,7 @@ from pathlib import Path import numpy as np -from .alpha_post import clean_alpha, enforce_trimap +from .alpha_post import clean_alpha, enforce_trimap, suppress_alpha_by_chroma from .chroma import compute_bg_confidence from .despill import despill_green from .foreground import estimate_foreground_rgb @@ -71,6 +71,7 @@ def _run_image( 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 = 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 d3601ba..9fb2c30 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -33,6 +33,12 @@ class AlphaPostSettings: fill_hole_area_ratio: float = 0.00002 alpha_floor: float = 0.002 alpha_ceil: float = 0.998 + # Chroma-guided alpha suppression: in the unknown band, pull alpha toward 0 + # where background (green) confidence is high, clearing spill in gaps/holes. + chroma_suppress: bool = True + chroma_suppress_bg_low: float = 0.35 + chroma_suppress_bg_high: float = 0.80 + chroma_suppress_strength: float = 1.0 @dataclass(frozen=True) diff --git a/configs/default.yaml b/configs/default.yaml index c187e37..2ec8849 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -10,6 +10,12 @@ trimap: unknown_radius_ratio: 0.012 fg_safe_radius_ratio: 0.006 +alpha_post: + chroma_suppress: true + chroma_suppress_bg_low: 0.35 + chroma_suppress_bg_high: 0.80 + chroma_suppress_strength: 1.0 + foreground: enabled: true method: ml