diff --git a/bgfilter/config.py b/bgfilter/config.py index 377c7f3..2414655 100644 --- a/bgfilter/config.py +++ b/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), ) diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 298d50f..7b9ff8f 100644 --- a/bgfilter/pipeline.py +++ b/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) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 81c9945..2298f8b 100644 --- a/bgfilter/settings.py +++ b/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() diff --git a/bgfilter/trimap.py b/bgfilter/trimap.py index 082e9f1..08369c1 100644 --- a/bgfilter/trimap.py +++ b/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 diff --git a/configs/default.yaml b/configs/default.yaml index 2ec8849..95571f4 100644 --- a/configs/default.yaml +++ b/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