Codex
99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
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,
|
|
bg_confidence: 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)
|
|
|
|
cv2 = require_cv2()
|
|
edge = (alpha > settings.edge_low) & (alpha < settings.edge_high)
|
|
if settings.edge_expand_radius > 0:
|
|
radius = settings.edge_expand_radius
|
|
kernel = cv2.getStructuringElement(
|
|
cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1)
|
|
)
|
|
edge = cv2.dilate(edge.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)
|
|
alpha_weight = np.maximum(
|
|
np.clip((1.0 - np.abs(alpha - 0.5) * 2.0), 0.0, 1.0),
|
|
settings.alpha_weight_floor,
|
|
)
|
|
bg_weight = (
|
|
settings.bg_confidence_weight * np.clip(bg_confidence, 0.0, 1.0)
|
|
+ (1.0 - settings.bg_confidence_weight)
|
|
)
|
|
green_weight = np.clip(excess / max(settings.green_excess_margin * 8.0, 1e-6), 0.0, 1.0)
|
|
mask = edge.astype(np.float32) * alpha_weight * bg_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)
|
|
|
|
# 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)
|