lhk229
2b3b62ed3f
The veto's suspect zone was defined by background-hue projection, so --no-chroma silently skipped it. Without a key colour there is no hue test to lean on, so in complex mode the gate instead requires the second opinion itself to be confidently near-empty (second_lo -> second_hi ramp, default 0.15 -> 0.40): regions the HR-matting model decisively rejects can be cleared, while thin strands it merely blurs to mid-alpha are untouched -- protecting exactly the crisp-strand advantage the pipeline has over a raw BiRefNet mask. Flat mode's gate is unchanged. Validation (BG_IMAGE01-04, cuda fp32): flat default path bit-identical; complex-mode veto touches 0.007-0.025% of pixels, visibly clearing blurred residue near strands and milky specks in hair gaps with no strand erosion. Costs the HR-matting forward in --no-chroma runs (0.61 -> 1.6 s/image warm GPU); disable with --no-cross-check. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
276 lines
10 KiB
Python
276 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
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, SegmentationSettings
|
|
from .trimap import (
|
|
fuse_trimap,
|
|
fuse_trimap_directional,
|
|
generate_trimap,
|
|
trimap_to_alpha_seed,
|
|
)
|
|
from .vitmatte_infer import ViTMatteRunner
|
|
|
|
|
|
@dataclass
|
|
class MattingResult:
|
|
"""In-memory result of one matting run.
|
|
|
|
``rgb`` is the despilled foreground (uint8 HxWx3) and ``alpha`` is float32 HxW in
|
|
0..1; together they compose the output RGBA. The remaining arrays are the
|
|
intermediate maps used for debug previews.
|
|
"""
|
|
|
|
rgb: np.ndarray
|
|
alpha: np.ndarray
|
|
metadata: dict
|
|
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
|
|
|
|
|
|
class MattingPipeline:
|
|
def __init__(
|
|
self,
|
|
settings: PipelineSettings,
|
|
*,
|
|
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:
|
|
from .segmentation import make_segmenter
|
|
|
|
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,
|
|
precision=check.precision,
|
|
)
|
|
)
|
|
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"
|
|
if self.settings.model.matting_method != "vitmatte":
|
|
raise RuntimeError(
|
|
f"Unsupported matting method '{self.settings.model.matting_method}'. "
|
|
"Use 'vitmatte' or 'chroma'."
|
|
)
|
|
|
|
try:
|
|
if self._runner is None:
|
|
self._runner = ViTMatteRunner(self.settings.model)
|
|
return self._runner.predict(rgb, trimap), "vitmatte"
|
|
except RuntimeError:
|
|
if not self.settings.model.fallback_to_chroma_alpha:
|
|
raise
|
|
return trimap_to_alpha_seed(trimap, bg_confidence), "chroma_fallback"
|
|
|
|
def run_rgb(self, rgb: np.ndarray) -> MattingResult:
|
|
"""Run the full pipeline on an in-memory RGB array (no disk I/O)."""
|
|
return _process_rgb(rgb, self)
|
|
|
|
def run_image(
|
|
self,
|
|
input_path: str | Path,
|
|
output_path: str | Path,
|
|
debug_dir: str | Path | None,
|
|
) -> dict:
|
|
return _run_image(input_path, output_path, debug_dir, self)
|
|
|
|
|
|
def run_image(
|
|
input_path: str | Path,
|
|
output_path: str | Path,
|
|
debug_dir: str | Path | None,
|
|
settings: PipelineSettings,
|
|
) -> dict:
|
|
return MattingPipeline(settings).run_image(input_path, output_path, debug_dir)
|
|
|
|
|
|
def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult:
|
|
settings = pipeline.settings
|
|
if settings.chroma.enabled:
|
|
screen_color = parse_hex_color(settings.screen_color) if settings.screen_color else None
|
|
bg_confidence, model = compute_bg_confidence(
|
|
rgb, settings=settings.chroma, screen_color=screen_color
|
|
)
|
|
else:
|
|
if not settings.segmentation.enabled:
|
|
raise RuntimeError(
|
|
"chroma.enabled: false (non-flat background mode) needs the segmentation "
|
|
"pipeline; enable segmentation or re-enable chroma."
|
|
)
|
|
if settings.model.matting_method == "chroma":
|
|
raise RuntimeError(
|
|
"matting_method 'chroma' needs the chroma key; use 'vitmatte' or re-enable chroma."
|
|
)
|
|
# Non-flat background: no colour key exists. A zero background-confidence
|
|
# map makes every chroma-fused formula degrade to its seg-only form; the
|
|
# colour-keyed stages (hue split, suppression, veto, despill) are skipped.
|
|
bg_confidence = np.zeros(rgb.shape[:2], dtype=np.float32)
|
|
model = None
|
|
second_alpha = None
|
|
if settings.segmentation.enabled:
|
|
reuse = (
|
|
settings.cross_check.enabled
|
|
and settings.cross_check.reuse_as_seg
|
|
# Reuse swaps one BiRefNet-family model for another (validated
|
|
# equivalent topology). A different backend chosen explicitly
|
|
# (anime-seg) keeps its own forward; cross-check still runs
|
|
# independently on top of it.
|
|
and settings.segmentation.backend == "birefnet"
|
|
)
|
|
if reuse:
|
|
# One HR-matting forward serves both the trimap topology and the
|
|
# cross-check second opinion; the primary seg model is never loaded.
|
|
second_alpha = pipeline._second_opinion(rgb)
|
|
seg_mask = second_alpha
|
|
else:
|
|
seg_mask = pipeline._segment(rgb)
|
|
# The directional modes are colour tests against the key colour; without
|
|
# one, the seg-topology trimap is the only meaningful choice.
|
|
mode = settings.trimap.mode if settings.chroma.enabled else "seg"
|
|
if mode == "seg":
|
|
trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
|
|
elif mode in ("directional", "directional-hard-bg"):
|
|
lab = convert_color_spaces(rgb)[1]
|
|
trimap, trimap_stats = fuse_trimap_directional(
|
|
seg_mask, bg_confidence, lab, model.lab_center, settings.trimap
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
f"Unknown trimap.mode '{mode}'. "
|
|
"Use 'directional', 'seg', or 'directional-hard-bg'."
|
|
)
|
|
else:
|
|
seg_mask = None
|
|
trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
|
|
raw_alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence)
|
|
|
|
alpha = enforce_trimap(raw_alpha, trimap)
|
|
if settings.chroma.enabled:
|
|
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)
|
|
if settings.cross_check.enabled:
|
|
lab = convert_color_spaces(rgb)[1]
|
|
if second_alpha is None:
|
|
second_alpha = pipeline._second_opinion(rgb)
|
|
alpha = cross_check_alpha(
|
|
alpha,
|
|
second_alpha,
|
|
# No key colour in complex-background mode: proj=None switches the
|
|
# veto to its second-opinion-confidence gate (see cross_check_alpha).
|
|
bg_hue_projection(lab, model.lab_center) if model is not None else None,
|
|
lab[..., 0],
|
|
trimap,
|
|
settings.cross_check,
|
|
)
|
|
foreground = estimate_foreground_rgb(
|
|
rgb, alpha, bg_confidence, model, settings.foreground
|
|
)
|
|
corrected_rgb, color_mask = despill(foreground, alpha, model, settings.despill)
|
|
|
|
metadata = {
|
|
"alpha_source": alpha_source,
|
|
"background_model": model.to_dict() if model is not None else None,
|
|
"trimap": trimap_stats,
|
|
"settings": asdict(settings),
|
|
}
|
|
return MattingResult(
|
|
rgb=corrected_rgb,
|
|
alpha=alpha,
|
|
metadata=metadata,
|
|
bg_confidence=bg_confidence,
|
|
trimap=trimap,
|
|
seg_mask=seg_mask,
|
|
cross_check_alpha=second_alpha,
|
|
color_mask=color_mask,
|
|
input_rgb=rgb,
|
|
)
|
|
|
|
|
|
def _run_image(
|
|
input_path: str | Path,
|
|
output_path: str | Path,
|
|
debug_dir: str | Path | None,
|
|
pipeline: MattingPipeline,
|
|
) -> dict:
|
|
res = _process_rgb(load_rgb(input_path), pipeline)
|
|
save_rgba(output_path, res.rgb, res.alpha)
|
|
|
|
result = {"input": str(input_path), "output": str(output_path), **res.metadata}
|
|
|
|
if debug_dir is not None:
|
|
debug = Path(debug_dir)
|
|
debug.mkdir(parents=True, exist_ok=True)
|
|
save_gray(debug / "bg_confidence.png", res.bg_confidence)
|
|
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)
|
|
save_previews(debug, res.rgb, res.alpha)
|
|
tiles = {
|
|
"input": res.input_rgb,
|
|
"bg confidence": np.repeat((res.bg_confidence[..., None] * 255).astype(np.uint8), 3, axis=2),
|
|
"trimap": np.repeat(res.trimap[..., None], 3, axis=2),
|
|
"alpha": np.repeat((res.alpha[..., None] * 255).astype(np.uint8), 3, axis=2),
|
|
}
|
|
for name, color in BACKGROUND_COLORS.items():
|
|
tiles[f"preview {name}"] = composite(res.rgb, res.alpha, color)
|
|
make_qa_grid(debug, tiles)
|
|
write_text(debug / "metadata.json", json.dumps(result, indent=2, ensure_ascii=False))
|
|
|
|
return result
|