Implement green screen matting pipeline

bgfilter/trimap.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import numpy as np
+
+from .deps import require_cv2
+from .settings import TrimapSettings
+
+
+def radius_from_ratio(shape: tuple[int, int], ratio: float, minimum: int) -> int:
+    return max(minimum, int(round(max(shape) * ratio)))
+
+
+def elliptical_kernel(radius: int) -> np.ndarray:
+    cv2 = require_cv2()
+    size = radius * 2 + 1
+    return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size))
+
+
+def generate_trimap(
+    bg_confidence: np.ndarray, settings: TrimapSettings
+) -> tuple[np.ndarray, dict[str, int]]:
+    cv2 = require_cv2()
+    shape = bg_confidence.shape
+    unknown_radius = radius_from_ratio(
+        shape, settings.unknown_radius_ratio, settings.min_unknown_radius
+    )
+    fg_safe_radius = radius_from_ratio(
+        shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius
+    )
+
+    sure_bg = bg_confidence >= settings.sure_bg_threshold
+    low_bg = bg_confidence <= settings.sure_fg_threshold
+
+    bg_u8 = sure_bg.astype(np.uint8)
+    unknown_band = cv2.dilate(bg_u8, elliptical_kernel(unknown_radius)).astype(bool)
+    fg_safe = ~cv2.dilate(bg_u8, elliptical_kernel(fg_safe_radius)).astype(bool)
+    sure_fg = low_bg & fg_safe
+
+    trimap = np.full(shape, 128, dtype=np.uint8)
+    trimap[sure_bg] = 0
+    trimap[sure_fg] = 255
+
+    # Keep a protective unknown band around all sure background, including holes.
+    trimap[unknown_band & ~sure_bg & ~sure_fg] = 128
+
+    stats = {
+        "sure_bg_pixels": int((trimap == 0).sum()),
+        "unknown_pixels": int((trimap == 128).sum()),
+        "sure_fg_pixels": int((trimap == 255).sum()),
+        "unknown_radius": int(unknown_radius),
+        "fg_safe_radius": int(fg_safe_radius),
+    }
+    return trimap, stats
+
+
+def trimap_to_alpha_seed(trimap: np.ndarray, bg_confidence: np.ndarray) -> np.ndarray:
+    alpha = np.clip(1.0 - bg_confidence, 0.0, 1.0).astype(np.float32)
+    alpha[trimap == 0] = 0.0
+    alpha[trimap == 255] = 1.0
+    return alpha