Add sample smoke check and refine foreground despill

bgfilter/despill.py

@@ -7,6 +7,18 @@ 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,
@@ -71,4 +83,16 @@ def despill_green(
     )
     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)

bgfilter/settings.py

@@ -47,6 +47,10 @@ class DespillSettings:
     alpha_weight_floor: float = 0.35
     unmix_strength: float = 0.55
     unmix_alpha_min: float = 0.08
+    local_fg_strength: float = 0.35
+    local_fg_blur_radius: int = 9
+    local_fg_alpha_threshold: float = 0.92
+    local_fg_bg_confidence_max: float = 0.20
 
 
 @dataclass(frozen=True)