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:
@@ -11,10 +11,14 @@ def despill_green(
|
||||
alpha: np.ndarray,
|
||||
model: BackgroundModel,
|
||||
settings: DespillSettings,
|
||||
screen_color: tuple[float, float, float] | None = None,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
rgb_f = rgb.astype(np.float32) / 255.0
|
||||
if not settings.enabled:
|
||||
return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
|
||||
if screen_color is not None:
|
||||
# Arbitrary background colour: despill along its Lab chroma direction.
|
||||
return _despill_color(rgb, alpha, model, settings)
|
||||
|
||||
# (1) Despill across the whole foreground, not only the soft edge band.
|
||||
# Green frequently bleeds into pixels ViTMatte marks as solid foreground
|
||||
@@ -65,3 +69,58 @@ def despill_green(
|
||||
out[..., 2] = np.clip(out[..., 2] + compensation * 0.5, 0.0, 1.0)
|
||||
|
||||
return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)
|
||||
|
||||
|
||||
def _despill_color(
|
||||
rgb: np.ndarray,
|
||||
alpha: np.ndarray,
|
||||
model: BackgroundModel,
|
||||
settings: DespillSettings,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Despill toward an arbitrary background colour.
|
||||
|
||||
Generalises the green channel algebra: the background's Lab a*/b* vector gives
|
||||
a chroma direction; foreground chroma projecting along it beyond a margin is
|
||||
treated as spill and pulled back, keeping luminance so texture survives. Reduces
|
||||
to a green-style suppression when the background is a saturated primary, but also
|
||||
handles low-chroma pastels where no single RGB channel dominates.
|
||||
"""
|
||||
from .deps import require_cv2
|
||||
|
||||
cv2 = require_cv2()
|
||||
rgb_f = rgb.astype(np.float32) / 255.0
|
||||
|
||||
bg_lab = np.asarray(model.lab_center, dtype=np.float32)
|
||||
va, vb = float(bg_lab[1]), float(bg_lab[2])
|
||||
mag = float(np.hypot(va, vb))
|
||||
if mag < 1e-3:
|
||||
return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
|
||||
da, db = va / mag, vb / mag
|
||||
|
||||
foreground = alpha > settings.edge_low
|
||||
if settings.edge_expand_radius > 0:
|
||||
radius = settings.edge_expand_radius
|
||||
kernel = cv2.getStructuringElement(
|
||||
cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1)
|
||||
)
|
||||
foreground = cv2.dilate(foreground.astype(np.uint8), kernel).astype(bool)
|
||||
|
||||
lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB)
|
||||
A, B = lab[..., 1], lab[..., 2]
|
||||
proj = A * da + B * db # chroma toward the background hue
|
||||
excess = np.maximum(proj - settings.color_margin_lab, 0.0)
|
||||
|
||||
alpha_weight = np.maximum(
|
||||
np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor
|
||||
)
|
||||
color_weight = np.clip(
|
||||
excess / max(settings.color_margin_lab * 8.0, 1e-6), 0.0, 1.0
|
||||
)
|
||||
mask = foreground.astype(np.float32) * alpha_weight * color_weight
|
||||
reduce = excess * mask * settings.strength # pull chroma back along the bg axis
|
||||
|
||||
lab_out = lab.copy()
|
||||
lab_out[..., 1] = A - reduce * da
|
||||
lab_out[..., 2] = B - reduce * db
|
||||
out = cv2.cvtColor(lab_out, cv2.COLOR_LAB2RGB)
|
||||
return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)
|
||||
|
||||
Reference in New Issue
Block a user