Simplify to single-segmenter pipelines; add screen_color prior; drop rescue + recolor
Strategic pivot: the dropout-fill rescue, the anime-seg fill backend, the
BiRefNet+anime-seg intersection, and the recolor pass all existed to repair
source images whose hair was already green-contaminated or broken at generation
time. Fixing the source instead (clean pastel-background generation) makes the
matte high-contrast, so that whole compensation layer is unnecessary. Remove it.
- Two pipelines via segmentation.enabled:
false -> chroma-only (Chroma + ViTMatte)
true -> single segmenter (default anime-seg, switchable birefnet) -> trimap
then ViTMatte refines, pymatting estimates foreground, despill cleans spill.
- Generalise the green-hardcoded colour logic to a screen_color prior (hex, e.g.
"#CFEFFF"; default null = green auto-detect, byte-identical). chroma keys off
Lab/RGB distance to the colour; despill removes chroma along its Lab direction.
Added --screen-color CLI flag.
- Removed dead code: fill_seg_dropouts + seg_fill*/intersection params,
SegmentationSettings.fill_backend/fill_model_name/sharpen, unsharp_mask,
recolor.py + RecolorSettings.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+8
-44
@@ -6,16 +6,10 @@ from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .alpha_post import (
|
||||
clean_alpha,
|
||||
enforce_trimap,
|
||||
fill_seg_dropouts,
|
||||
suppress_alpha_by_chroma,
|
||||
)
|
||||
from .chroma import compute_bg_confidence
|
||||
from .alpha_post import clean_alpha, enforce_trimap, suppress_alpha_by_chroma
|
||||
from .chroma import compute_bg_confidence, parse_hex_color
|
||||
from .despill import despill_green
|
||||
from .foreground import estimate_foreground_rgb
|
||||
from .recolor import recolor_foreground
|
||||
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
|
||||
@@ -28,36 +22,13 @@ class MattingPipeline:
|
||||
self.settings = settings
|
||||
self._runner: ViTMatteRunner | None = None
|
||||
self._segmenter = None
|
||||
self._fill_segmenter = None
|
||||
|
||||
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)
|
||||
mask = self._segmenter.mask(rgb)
|
||||
seg = self.settings.segmentation
|
||||
if seg.sharpen > 0:
|
||||
from .segmentation import unsharp_mask
|
||||
|
||||
mask = unsharp_mask(mask, seg.sharpen, seg.sharpen_radius)
|
||||
return mask
|
||||
|
||||
def _segment_fill(self, rgb: np.ndarray, default_mask: np.ndarray) -> np.ndarray:
|
||||
"""Mask for the dropout-fill; may use a different (higher-recall) backend
|
||||
than the trimap, since the fill only adds and cannot break anything."""
|
||||
seg = self.settings.segmentation
|
||||
if not seg.fill_backend or seg.fill_backend == seg.backend:
|
||||
return default_mask
|
||||
if self._fill_segmenter is None:
|
||||
from dataclasses import replace
|
||||
|
||||
from .segmentation import make_segmenter
|
||||
|
||||
self._fill_segmenter = make_segmenter(
|
||||
replace(seg, backend=seg.fill_backend, model_name=seg.fill_model_name)
|
||||
)
|
||||
return self._fill_segmenter.mask(rgb)
|
||||
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":
|
||||
@@ -103,7 +74,10 @@ def _run_image(
|
||||
) -> dict:
|
||||
settings = pipeline.settings
|
||||
rgb = load_rgb(input_path)
|
||||
bg_confidence, model = compute_bg_confidence(rgb, settings=settings.chroma)
|
||||
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
|
||||
)
|
||||
if settings.segmentation.enabled:
|
||||
seg_mask = pipeline._segment(rgb)
|
||||
trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
|
||||
@@ -115,22 +89,12 @@ def _run_image(
|
||||
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)
|
||||
if seg_mask is not None:
|
||||
fill_mask = pipeline._segment_fill(rgb, seg_mask)
|
||||
# seg_mask is the reliable trimap backend (BiRefNet); it gates the fill
|
||||
# when seg_fill_intersect is on (fill_mask ∧ seg_mask).
|
||||
alpha = fill_seg_dropouts(
|
||||
alpha, fill_mask, settings.alpha_post, gate_mask=seg_mask
|
||||
)
|
||||
foreground = estimate_foreground_rgb(
|
||||
rgb, alpha, bg_confidence, model, settings.foreground
|
||||
)
|
||||
corrected_rgb, color_mask = despill_green(
|
||||
foreground.rgb, alpha, model, settings.despill
|
||||
foreground.rgb, alpha, model, settings.despill, screen_color=screen_color
|
||||
)
|
||||
if settings.recolor.enabled:
|
||||
recolored = recolor_foreground(corrected_rgb, alpha, seg_mask, settings.recolor)
|
||||
corrected_rgb, color_mask = recolored.rgb, recolored.mask
|
||||
save_rgba(output_path, corrected_rgb, alpha)
|
||||
|
||||
result = {
|
||||
|
||||
Reference in New Issue
Block a user