Implement green screen matting pipeline

2026-06-30 14:14:15 +08:00
parent 4d2a2015fe
commit 8fc53925e0
17 changed files with 896 additions and 0 deletions
@@ -0,0 +1,39 @@
+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,
+    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)
+
+    edge = (alpha > settings.edge_low) & (alpha < settings.edge_high)
+    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)
+    edge_weight = np.clip((1.0 - np.abs(alpha - 0.5) * 2.0), 0.0, 1.0)
+    bg_weight = np.clip(bg_confidence, 0.0, 1.0)
+    mask = edge.astype(np.float32) * edge_weight * np.maximum(bg_weight, 0.25)
+
+    out = rgb_f.copy()
+    out[..., 1] = g - excess * mask * settings.strength
+
+    # Very light compensation toward the non-green channels to avoid gray fringes.
+    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)