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>
This commit is contained in:
2026-06-30 21:32:38 +08:00
parent a6a313d2d0
commit 711669ea51
5 changed files with 84 additions and 3 deletions
+17 -2
View File
@@ -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)