lhk229
9b7a192599
After the full-foreground despill (5ddf0cd) the residual green was confined to
semi-transparent hair (alpha < 0.5): 85% of the worst pixels had alpha < 0.3,
where a pixel-wise unmix (divide by small alpha) is too noisy to trust.
Replace the hand-rolled unmix/local-blur estimator with pymatting's
estimate_foreground_ml (Germer et al. multi-level closed form), which propagates
reliable foreground colour from high-alpha neighbours into the fringe and
estimates the background, so green spill is unmixed rather than clamped. The
despill pass stays as a light cleanup on top.
- foreground.method selects "ml" (default) or "unmix" (legacy heuristic kept as
a fallback when pymatting is unavailable, matching the chroma fallback idiom).
- ForegroundEstimate now exposes rgb (F), background (B) and a correction map;
debug outputs become foreground_rgb / foreground_background / foreground_correction.
Controlled comparison (same ViTMatte alpha, only foreground method changed),
edge_green_excess:
TestImage p95 0.176->0.031, mean 0.061->0.005
TestImage2 p95 0.165->0.004, mean 0.052->0.002
Adds pymatting (pulls in numba) to requirements.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from .alpha_post import clean_alpha, enforce_trimap
|
|
from .chroma import compute_bg_confidence
|
|
from .despill import despill_green
|
|
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 .vitmatte_infer import ViTMatteRunner
|
|
|
|
|
|
class MattingPipeline:
|
|
def __init__(self, settings: PipelineSettings):
|
|
self.settings = settings
|
|
self._runner: ViTMatteRunner | None = None
|
|
|
|
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_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 _run_image(
|
|
input_path: str | Path,
|
|
output_path: str | Path,
|
|
debug_dir: str | Path | None,
|
|
pipeline: MattingPipeline,
|
|
) -> dict:
|
|
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)
|
|
alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence)
|
|
|
|
alpha = enforce_trimap(alpha, trimap)
|
|
alpha = clean_alpha(alpha, trimap, settings.alpha_post)
|
|
foreground = estimate_foreground_rgb(
|
|
rgb, alpha, bg_confidence, model, settings.foreground
|
|
)
|
|
corrected_rgb, despill_mask = despill_green(
|
|
foreground.rgb, alpha, model, settings.despill
|
|
)
|
|
save_rgba(output_path, corrected_rgb, alpha)
|
|
|
|
result = {
|
|
"input": str(input_path),
|
|
"output": str(output_path),
|
|
"alpha_source": alpha_source,
|
|
"background_model": model.to_dict(),
|
|
"trimap": trimap_stats,
|
|
"settings": asdict(settings),
|
|
}
|
|
|
|
if debug_dir is not None:
|
|
debug = Path(debug_dir)
|
|
debug.mkdir(parents=True, exist_ok=True)
|
|
save_gray(debug / "bg_confidence.png", bg_confidence)
|
|
save_gray(debug / "trimap.png", trimap)
|
|
save_gray(debug / "alpha.png", alpha)
|
|
save_rgb(debug / "foreground_rgb.png", foreground.rgb)
|
|
save_rgb(debug / "foreground_background.png", foreground.background)
|
|
save_gray(debug / "foreground_correction.png", foreground.correction)
|
|
save_gray(debug / "despill_mask.png", despill_mask)
|
|
save_previews(debug, corrected_rgb, alpha)
|
|
tiles = {
|
|
"input": rgb,
|
|
"bg confidence": np.repeat((bg_confidence[..., None] * 255).astype(np.uint8), 3, axis=2),
|
|
"trimap": np.repeat(trimap[..., None], 3, axis=2),
|
|
"alpha": np.repeat((alpha[..., None] * 255).astype(np.uint8), 3, axis=2),
|
|
}
|
|
for name, color in BACKGROUND_COLORS.items():
|
|
tiles[f"preview {name}"] = composite(corrected_rgb, alpha, color)
|
|
make_qa_grid(debug, tiles)
|
|
write_text(debug / "metadata.json", json.dumps(result, indent=2, ensure_ascii=False))
|
|
|
|
return result
|