Fuse semantic segmentation mask into the trimap

Wire BiRefNet into the pipeline (lazy-loaded) and build the trimap from the
semantic subject mask refined by the green key, used when segmentation is
enabled. generate_trimap (chroma-only) stays as the no-segmentation fallback.

Authority split: the mask decides subject topology (keeps green-tinted hair as
foreground, drops see-through holes), the bright-green key sharpens the screen
boundary, ViTMatte refines a small protective band. Background follows the mask
directly so interior holes survive; the band is small so it does not swallow
finger gaps.

End-to-end on the samples: green hair preserved, finger-gap green 99%
transparent (BiRefNet marks it background; was 100% foreground before), leg
edges smooth -- topology/colour/edge problems local cues could not solve, with
no hard alpha surgery.

- SegmentationSettings wired into PipelineSettings/config; --device propagates.
- TrimapSettings: seg_core_threshold, seg_loose_threshold.
- New debug output seg_mask.png.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

bgfilter/config.py

@@ -11,6 +11,7 @@ from .settings import (
     ForegroundSettings,
     ModelSettings,
     PipelineSettings,
+    SegmentationSettings,
     TrimapSettings,
 )
 
@@ -42,7 +43,7 @@ def _update_dataclass(instance: T, values: dict[str, Any] | None) -> T:
 
 
 def settings_from_dict(data: dict[str, Any]) -> PipelineSettings:
-    allowed_sections = {"chroma", "trimap", "alpha_post", "foreground", "despill", "model"}
+    allowed_sections = {"chroma", "trimap", "alpha_post", "foreground", "despill", "model", "segmentation"}
     unknown_sections = sorted(set(data) - allowed_sections)
     if unknown_sections:
         raise ValueError(f"Unknown config section(s): {', '.join(unknown_sections)}")
@@ -54,6 +55,7 @@ def settings_from_dict(data: dict[str, Any]) -> PipelineSettings:
         foreground=_update_dataclass(ForegroundSettings(), data.get("foreground")),
         despill=_update_dataclass(DespillSettings(), data.get("despill")),
         model=_update_dataclass(ModelSettings(), data.get("model")),
+        segmentation=_update_dataclass(SegmentationSettings(), data.get("segmentation")),
     )
 
 
@@ -70,6 +72,7 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS
     foreground = settings.foreground
     despill = settings.despill
     model = settings.model
+    segmentation = settings.segmentation
 
     model_updates = {
         key: overrides[key]
@@ -89,6 +92,9 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS
     despill_updates: dict[str, Any] = {}
     if overrides.get("despill") is not None:
         despill_updates["enabled"] = overrides["despill"]
+    seg_updates: dict[str, Any] = {}
+    if overrides.get("device") is not None:
+        seg_updates["device"] = overrides["device"]
 
     return PipelineSettings(
         chroma=chroma,
@@ -97,4 +103,5 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS
         foreground=foreground,
         despill=_update_dataclass(despill, despill_updates),
         model=_update_dataclass(model, model_updates),
+        segmentation=_update_dataclass(segmentation, seg_updates),
     )

bgfilter/pipeline.py

@@ -13,7 +13,7 @@ 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 generate_trimap, trimap_to_alpha_seed
+from .trimap import fuse_trimap, generate_trimap, trimap_to_alpha_seed
 from .vitmatte_infer import ViTMatteRunner
 
 
@@ -21,6 +21,14 @@ class MattingPipeline:
     def __init__(self, settings: PipelineSettings):
         self.settings = settings
         self._runner: ViTMatteRunner | None = None
+        self._segmenter = None
+
+    def _segment(self, rgb: np.ndarray) -> np.ndarray:
+        if self._segmenter is None:
+            from .segmentation import BiRefNetSegmenter
+
+            self._segmenter = BiRefNetSegmenter(self.settings.segmentation)
+        return self._segmenter.mask(rgb)
 
     def _predict_alpha(self, rgb: np.ndarray, trimap: np.ndarray, bg_confidence: np.ndarray) -> tuple[np.ndarray, str]:
         if self.settings.model.matting_method == "chroma":
@@ -67,7 +75,12 @@ def _run_image(
     settings = pipeline.settings
     rgb = load_rgb(input_path)
     bg_confidence, model = compute_bg_confidence(rgb, settings=settings.chroma)
-    trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
+    if settings.segmentation.enabled:
+        seg_mask = pipeline._segment(rgb)
+        trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
+    else:
+        seg_mask = None
+        trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
     alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence)
 
     alpha = enforce_trimap(alpha, trimap)
@@ -95,6 +108,8 @@ def _run_image(
         debug.mkdir(parents=True, exist_ok=True)
         save_gray(debug / "bg_confidence.png", bg_confidence)
         save_gray(debug / "trimap.png", trimap)
+        if seg_mask is not None:
+            save_gray(debug / "seg_mask.png", seg_mask)
         save_gray(debug / "alpha.png", alpha)
         save_rgb(debug / "foreground_rgb.png", foreground.rgb)
         save_rgb(debug / "foreground_background.png", foreground.background)

bgfilter/settings.py

@@ -25,6 +25,9 @@ class TrimapSettings:
     fg_safe_radius_ratio: float = 0.006
     min_unknown_radius: int = 4
     min_fg_safe_radius: int = 2
+    # Semantic fusion thresholds (used by fuse_trimap when segmentation is on).
+    seg_core_threshold: float = 0.60
+    seg_loose_threshold: float = 0.15
 
 
 @dataclass(frozen=True)
@@ -94,3 +97,4 @@ class PipelineSettings:
     foreground: ForegroundSettings = ForegroundSettings()
     despill: DespillSettings = DespillSettings()
     model: ModelSettings = ModelSettings()
+    segmentation: SegmentationSettings = SegmentationSettings()

bgfilter/trimap.py

@@ -53,6 +53,53 @@ def generate_trimap(
     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 green key.
+
+    Authority split: the segmentation mask decides subject *topology* (it keeps
+    colour-contaminated hair as foreground and drops see-through holes), while the
+    bright-green chroma key sharpens the screen 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)
+
+    green_bg = bg_confidence >= settings.sure_bg_threshold       # bright clean screen
+    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 bright screen.
+    sure_fg = core.astype(bool) & ~band & ~green_bg
+    sure_bg = ((loose == 0) & ~band) | (green_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 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

configs/default.yaml

@@ -4,11 +4,19 @@ model:
   matting_method: vitmatte
   fallback_to_chroma_alpha: false
 
+segmentation:
+  enabled: true
+  model_name: ZhengPeng7/BiRefNet
+  device: cuda
+  input_size: 1024
+
 trimap:
   sure_bg_threshold: 0.92
   sure_fg_threshold: 0.12
   unknown_radius_ratio: 0.012
   fg_safe_radius_ratio: 0.006
+  seg_core_threshold: 0.60
+  seg_loose_threshold: 0.15
 
 alpha_post:
   chroma_suppress: true