Merge master: matte-confidence gate + cross-model hair-gap veto
Brings in the two algorithm optimisations from master: - Gate chroma suppression by the matte's own confidence (639fdf0) - Add cross-model veto of hair-gap background residue (adbbb9d) Server-edition integration on top of the plain merge: - PipelineManager caches the cross-check model (BiRefNet_HR-matting) and injects it into MattingPipeline, so HTTP requests never reload it. - MattingResult carries cross_check_alpha; debug output matches master. - /healthz reports crossCheck; README lists models/BiRefNet_HR-matting in the local-weights bundle. Verified: CLI --no-cross-check baseline and full cross-check run both pass on samples/TestImage3.png; cross_check_alpha.png written to debug. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+53
-8
@@ -6,13 +6,23 @@ from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .alpha_post import clean_alpha, enforce_trimap, suppress_alpha_by_chroma
|
||||
from .chroma import compute_bg_confidence, parse_hex_color
|
||||
from .alpha_post import (
|
||||
clean_alpha,
|
||||
cross_check_alpha,
|
||||
enforce_trimap,
|
||||
suppress_alpha_by_chroma,
|
||||
)
|
||||
from .chroma import (
|
||||
bg_hue_projection,
|
||||
compute_bg_confidence,
|
||||
convert_color_spaces,
|
||||
parse_hex_color,
|
||||
)
|
||||
from .despill import despill
|
||||
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 .settings import PipelineSettings, SegmentationSettings
|
||||
from .trimap import (
|
||||
fuse_trimap,
|
||||
fuse_trimap_directional,
|
||||
@@ -37,6 +47,7 @@ class MattingResult:
|
||||
bg_confidence: np.ndarray
|
||||
trimap: np.ndarray
|
||||
seg_mask: np.ndarray | None
|
||||
cross_check_alpha: np.ndarray | None
|
||||
color_mask: np.ndarray
|
||||
input_rgb: np.ndarray
|
||||
|
||||
@@ -48,12 +59,14 @@ class MattingPipeline:
|
||||
*,
|
||||
runner: ViTMatteRunner | None = None,
|
||||
segmenter=None,
|
||||
cross_checker=None,
|
||||
):
|
||||
self.settings = settings
|
||||
# Pre-loaded models can be injected so a long-lived service shares them
|
||||
# across requests instead of reloading per call; None = lazy-load on first use.
|
||||
self._runner: ViTMatteRunner | None = runner
|
||||
self._segmenter = segmenter
|
||||
self._cross_checker = cross_checker
|
||||
|
||||
def _segment(self, rgb: np.ndarray) -> np.ndarray:
|
||||
if self._segmenter is None:
|
||||
@@ -62,6 +75,20 @@ class MattingPipeline:
|
||||
self._segmenter = make_segmenter(self.settings.segmentation)
|
||||
return self._segmenter.mask(rgb)
|
||||
|
||||
def _second_opinion(self, rgb: np.ndarray) -> np.ndarray:
|
||||
if self._cross_checker is None:
|
||||
from .segmentation import BiRefNetSegmenter
|
||||
|
||||
check = self.settings.cross_check
|
||||
self._cross_checker = BiRefNetSegmenter(
|
||||
SegmentationSettings(
|
||||
model_name=check.model_name,
|
||||
device=self.settings.segmentation.device,
|
||||
input_size=check.input_size,
|
||||
)
|
||||
)
|
||||
return self._cross_checker.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":
|
||||
return trimap_to_alpha_seed(trimap, bg_confidence), "chroma"
|
||||
@@ -114,8 +141,6 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult:
|
||||
if mode == "seg":
|
||||
trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
|
||||
elif mode in ("directional", "directional-hard-bg"):
|
||||
from .chroma import convert_color_spaces
|
||||
|
||||
lab = convert_color_spaces(rgb)[1]
|
||||
trimap, trimap_stats = fuse_trimap_directional(
|
||||
seg_mask, bg_confidence, lab, model.lab_center, settings.trimap
|
||||
@@ -128,11 +153,28 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult:
|
||||
else:
|
||||
seg_mask = None
|
||||
trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
|
||||
alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence)
|
||||
raw_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 = enforce_trimap(raw_alpha, trimap)
|
||||
alpha = suppress_alpha_by_chroma(
|
||||
alpha, bg_confidence, trimap, settings.alpha_post,
|
||||
# The matte-confidence gate only makes sense for a real matting prediction;
|
||||
# a chroma-seeded alpha is itself colour evidence, so no gate there.
|
||||
raw_alpha=raw_alpha if alpha_source == "vitmatte" else None,
|
||||
)
|
||||
alpha = clean_alpha(alpha, trimap, settings.alpha_post)
|
||||
second_alpha = None
|
||||
if settings.cross_check.enabled:
|
||||
lab = convert_color_spaces(rgb)[1]
|
||||
second_alpha = pipeline._second_opinion(rgb)
|
||||
alpha = cross_check_alpha(
|
||||
alpha,
|
||||
second_alpha,
|
||||
bg_hue_projection(lab, model.lab_center),
|
||||
lab[..., 0],
|
||||
trimap,
|
||||
settings.cross_check,
|
||||
)
|
||||
foreground = estimate_foreground_rgb(
|
||||
rgb, alpha, bg_confidence, model, settings.foreground
|
||||
)
|
||||
@@ -151,6 +193,7 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult:
|
||||
bg_confidence=bg_confidence,
|
||||
trimap=trimap,
|
||||
seg_mask=seg_mask,
|
||||
cross_check_alpha=second_alpha,
|
||||
color_mask=color_mask,
|
||||
input_rgb=rgb,
|
||||
)
|
||||
@@ -174,6 +217,8 @@ def _run_image(
|
||||
save_gray(debug / "trimap.png", res.trimap)
|
||||
if res.seg_mask is not None:
|
||||
save_gray(debug / "seg_mask.png", res.seg_mask)
|
||||
if res.cross_check_alpha is not None:
|
||||
save_gray(debug / "cross_check_alpha.png", res.cross_check_alpha)
|
||||
save_gray(debug / "alpha.png", res.alpha)
|
||||
save_rgb(debug / "foreground_rgb.png", res.rgb)
|
||||
save_gray(debug / "color_mask.png", res.color_mask)
|
||||
|
||||
Reference in New Issue
Block a user