diff --git a/bgfilter/alpha_post.py b/bgfilter/alpha_post.py index 5151c32..72911fa 100644 --- a/bgfilter/alpha_post.py +++ b/bgfilter/alpha_post.py @@ -85,3 +85,24 @@ def clean_alpha( out[hole & (trimap == 255)] = 1.0 return enforce_trimap(out, trimap) + + +def fill_seg_dropouts( + alpha: np.ndarray, seg_mask: np.ndarray, settings: AlphaPostSettings +) -> np.ndarray: + """Rescue alpha where ViTMatte dropped a region the semantic mask asserts. + + Green-contaminated wisps are nearly the background colour, so ViTMatte (a + contrast-based matter) cuts them even though BiRefNet still sees them. Where + ViTMatte's alpha is near zero but the segmentation mask is confidently + foreground, restore a faint alpha from the mask. Limited to those dropouts so + ViTMatte's sharper edges elsewhere are untouched, and gated by seg_min so the + mask's faint halo and the finger gap (low mask) are not filled. + """ + if not settings.seg_fill: + return alpha + out = alpha.astype(np.float32).copy() + rescued = settings.seg_fill_strength * np.clip(seg_mask, 0.0, 1.0) + dropout = (out < settings.seg_fill_alpha_max) & (seg_mask >= settings.seg_fill_seg_min) + out[dropout] = np.maximum(out[dropout], rescued[dropout]) + return np.clip(out, 0.0, 1.0) diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 7b9ff8f..e346670 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -6,7 +6,12 @@ from pathlib import Path import numpy as np -from .alpha_post import clean_alpha, enforce_trimap, suppress_alpha_by_chroma +from .alpha_post import ( + clean_alpha, + enforce_trimap, + fill_seg_dropouts, + suppress_alpha_by_chroma, +) from .chroma import compute_bg_confidence from .despill import despill_green from .foreground import estimate_foreground_rgb @@ -86,6 +91,8 @@ def _run_image( 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) + if seg_mask is not None: + alpha = fill_seg_dropouts(alpha, seg_mask, 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 a9e1250..20e5508 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -42,6 +42,12 @@ class AlphaPostSettings: chroma_suppress_bg_low: float = 0.35 chroma_suppress_bg_high: float = 0.80 chroma_suppress_strength: float = 1.0 + # Segmentation dropout-fill: rescue alpha where ViTMatte cut a region the + # semantic mask asserts is foreground (green-contaminated wisps it cannot see). + seg_fill: bool = True + seg_fill_alpha_max: float = 0.10 + seg_fill_seg_min: float = 0.10 + seg_fill_strength: float = 0.70 @dataclass(frozen=True) diff --git a/configs/default.yaml b/configs/default.yaml index 13dc8f2..a1b3c47 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -23,6 +23,10 @@ alpha_post: chroma_suppress_bg_low: 0.35 chroma_suppress_bg_high: 0.80 chroma_suppress_strength: 1.0 + seg_fill: true + seg_fill_alpha_max: 0.10 + seg_fill_seg_min: 0.10 + seg_fill_strength: 0.70 foreground: enabled: true