lhk229
5ddf0cd392
Green spill bleeds into pixels ViTMatte marks as solid foreground (alpha~1), but despill was gated to the soft edge band (0.005<alpha<0.995) and weakened exactly there: the alpha tent peaked at 0.5 and bg_confidence (low on the character edge) further suppressed the correction. - (1) Apply despill to the whole foreground (alpha > edge_low) instead of the soft band only; neutral pixels stay untouched via the green-excess term. - (2) Weight by alpha at full strength across the opaque range (alpha>=0.5), easing off only where mostly transparent, instead of a tent peaking at 0.5. - (3) Drop bg_confidence as a despill gate; localise purely by green excess. Controlled comparison (same alpha, only despill changed), edge_green_excess: ViTMatte TestImage p95 0.239->0.176, mean 0.131->0.061 ViTMatte TestImage2 p95 0.231->0.165, mean 0.115->0.052 Removes now-unused DespillSettings fields edge_high and bg_confidence_weight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
68 lines
2.7 KiB
Python
68 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from .chroma import BackgroundModel
|
|
from .settings import DespillSettings
|
|
|
|
|
|
def despill_green(
|
|
rgb: np.ndarray,
|
|
alpha: np.ndarray,
|
|
model: BackgroundModel,
|
|
settings: DespillSettings,
|
|
) -> 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)
|
|
|
|
# (1) Despill across the whole foreground, not only the soft edge band.
|
|
# Green frequently bleeds into pixels ViTMatte marks as solid foreground
|
|
# (alpha ~ 1); the old upper bound left those uncorrected. Neutral pixels
|
|
# stay untouched because the green-excess term below is zero for them.
|
|
foreground = alpha > settings.edge_low
|
|
if settings.edge_expand_radius > 0:
|
|
from .deps import require_cv2
|
|
|
|
cv2 = require_cv2()
|
|
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)
|
|
|
|
r = rgb_f[..., 0]
|
|
g = rgb_f[..., 1]
|
|
b = rgb_f[..., 2]
|
|
neutral_green = np.maximum(r, b) + settings.green_excess_margin
|
|
excess = np.maximum(g - neutral_green, 0.0)
|
|
|
|
# (2) Keep full strength across the whole opaque range (alpha >= 0.5) and
|
|
# ease off only where the pixel is mostly transparent. The previous tent
|
|
# peaked at alpha 0.5 and starved exactly the solid hair that needs despill.
|
|
alpha_weight = np.maximum(
|
|
np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor
|
|
)
|
|
|
|
# (3) Localise purely by green excess. bg_confidence is intentionally not a
|
|
# gate: it is low on the character edge where the spill lives, so using it
|
|
# suppressed the very correction we want.
|
|
green_weight = np.clip(excess / max(settings.green_excess_margin * 8.0, 1e-6), 0.0, 1.0)
|
|
mask = foreground.astype(np.float32) * alpha_weight * green_weight
|
|
|
|
out = rgb_f.copy()
|
|
target_green = np.minimum(g, neutral_green)
|
|
out[..., 1] = g * (1.0 - mask * settings.strength) + target_green * (
|
|
mask * settings.strength
|
|
)
|
|
|
|
# Light foreground color approximation: green spill usually removed too much
|
|
# luminance from thin hair edges, so redistribute part of the removed green
|
|
# into red/blue without changing fully opaque pixels.
|
|
bg_green = float(model.rgb_center[1])
|
|
compensation = excess * mask * settings.strength * min(0.25, bg_green * 0.15)
|
|
out[..., 0] = np.clip(out[..., 0] + compensation * 0.5, 0.0, 1.0)
|
|
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)
|