Add foreground color estimation pass

This commit is contained in:
Codex
2026-06-30 16:25:56 +08:00
parent 1a670088c0
commit d6bb658244
9 changed files with 158 additions and 70 deletions
+3 -40
View File
@@ -3,22 +3,9 @@ from __future__ import annotations
import numpy as np
from .chroma import BackgroundModel
from .deps import require_cv2
from .settings import DespillSettings
def _weighted_blur(values: np.ndarray, weights: np.ndarray, radius: int) -> np.ndarray:
if radius <= 0:
return values
cv2 = require_cv2()
kernel_size = radius * 2 + 1
numerator = cv2.GaussianBlur(
values * weights[..., None], (kernel_size, kernel_size), 0
)
denominator = cv2.GaussianBlur(weights, (kernel_size, kernel_size), 0)
return numerator / np.maximum(denominator[..., None], 1e-6)
def despill_green(
rgb: np.ndarray,
alpha: np.ndarray,
@@ -30,9 +17,11 @@ def despill_green(
if not settings.enabled:
return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
cv2 = require_cv2()
edge = (alpha > settings.edge_low) & (alpha < settings.edge_high)
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)
@@ -69,30 +58,4 @@ def despill_green(
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)
# Approximate foreground recovery for semi-transparent contaminated pixels.
# C = alpha * F + (1 - alpha) * B, so F = (C - (1 - alpha) * B) / alpha.
# Blend it only where green excess is present; this is intentionally local
# because incorrect unmixing on opaque clothing would damage normal colors.
bg = np.asarray(model.rgb_center, dtype=np.float32)
a = np.clip(alpha[..., None], settings.unmix_alpha_min, 1.0)
recovered = np.clip((rgb_f - (1.0 - alpha[..., None]) * bg) / a, 0.0, 1.0)
unmix_mask = (
mask
* settings.unmix_strength
* np.clip((1.0 - alpha) / max(1.0 - settings.unmix_alpha_min, 1e-6), 0.0, 1.0)
)
out = out * (1.0 - unmix_mask[..., None]) + recovered * unmix_mask[..., None]
local_fg_weights = (
(alpha >= settings.local_fg_alpha_threshold)
& (bg_confidence <= settings.local_fg_bg_confidence_max)
).astype(np.float32)
local_fg = _weighted_blur(out, local_fg_weights, settings.local_fg_blur_radius)
local_mask = (
mask
* settings.local_fg_strength
* np.clip((1.0 - alpha) / max(1.0 - settings.unmix_alpha_min, 1e-6), 0.0, 1.0)
)
out = out * (1.0 - local_mask[..., None]) + local_fg * local_mask[..., None]
return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)