lhk229
f7e12d3498
The green in the finger gap survives keying because it is dark and desaturated (brightness ~0.58, hue-shifted), so chroma confidence reads ~0 and the trimap leaves it in the unknown band, where ViTMatte fills the pocket as foreground. It is also connected to the exterior screen (an open notch), so neither confidence suppression nor classic hole-filling removes it. Add carve_background_holes (doc section 10): a pixel that is clearly green-dominant in the input *and* mostly surrounded by genuine (non-green) foreground is set to alpha 0. The surround cue separates such an intrusion (finger gaps, inter-strand pockets) from the outer hair edge, which has background on one side and is preserved. Runs after alpha cleanup. On TestImage the finger-gap survivors (green & alpha>0.5) drop 170 -> 28 with no visible loss of hair wisps; inter-strand green pockets clear too. Knobs in AlphaPostSettings (hole_*, default on, surround_min 0.4). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
122 lines
4.4 KiB
Python
122 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from .alpha_post import (
|
|
carve_background_holes,
|
|
clean_alpha,
|
|
enforce_trimap,
|
|
suppress_alpha_by_chroma,
|
|
)
|
|
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 = suppress_alpha_by_chroma(alpha, bg_confidence, trimap, settings.alpha_post)
|
|
alpha = clean_alpha(alpha, trimap, settings.alpha_post)
|
|
alpha = carve_background_holes(alpha, rgb, 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
|