Files
lhk229 e61fbe84f0 Add non-flat background mode (--no-chroma / chroma.enabled: false)
No colour key at all: segmentation alone drives the trimap (mode forced
to 'seg', whose chroma terms degrade to no-ops on a zero background-
confidence map), ViTMatte still refines the unknown band at full
resolution, and pymatting still estimates edge foreground colour. Every
colour-keyed stage is bypassed: background auto-detect, the directional
hue split, chroma alpha suppression, despill, the cross-check veto (its
suspect zone is background-hued by definition; with reuse_as_seg the
second opinion still serves as the segmenter), and the unmix fallback.

Guards: requires segmentation.enabled and matting_method 'vitmatte'
(clear errors otherwise); despill/foreground handle model=None.

Validation: flat-background default path is bit-identical pre/post
(alpha and RGB |D|max = 0 on TestImage3). On a known-alpha subject
composited over a gradient+blotch background, recovered alpha scores
MAE 0.0018 / IoU@0.5 0.993 (0.995 with --cross-check-as-seg); the same
input correctly fails auto-detection in default mode.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-14 11:13:08 +08:00

62 lines
2.7 KiB
Python

from __future__ import annotations
import numpy as np
from .chroma import BackgroundModel
from .deps import require_cv2
from .settings import DespillSettings
def despill(
rgb: np.ndarray,
alpha: np.ndarray,
model: BackgroundModel | None,
settings: DespillSettings,
) -> tuple[np.ndarray, np.ndarray]:
"""Remove background-colour spill from the foreground.
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 toward
neutral, keeping luminance so hair texture and shading survive. This works for
any flat background: a saturated primary (green/blue screen) reduces to classic
channel suppression, while low-chroma pastels -- where no single RGB channel
dominates -- are handled by the same projection.
Returns the corrected RGB (uint8) and the per-pixel despill weight (float32).
"""
# model is None with chroma disabled: no key colour means no spill direction.
if not settings.enabled or model is None:
return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
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: # near-neutral background: no chroma direction, nothing to despill
return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
da, db = va / mag, vb / mag
cv2 = require_cv2()
rgb_f = rgb.astype(np.float32) / 255.0
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)
# Full strength across the opaque range; ease off only where mostly transparent.
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)