Use pymatting ML foreground estimation to remove edge spill

After the full-foreground despill (5ddf0cd) the residual green was confined to
semi-transparent hair (alpha < 0.5): 85% of the worst pixels had alpha < 0.3,
where a pixel-wise unmix (divide by small alpha) is too noisy to trust.

Replace the hand-rolled unmix/local-blur estimator with pymatting's
estimate_foreground_ml (Germer et al. multi-level closed form), which propagates
reliable foreground colour from high-alpha neighbours into the fringe and
estimates the background, so green spill is unmixed rather than clamped. The
despill pass stays as a light cleanup on top.

- foreground.method selects "ml" (default) or "unmix" (legacy heuristic kept as
  a fallback when pymatting is unavailable, matching the chroma fallback idiom).
- ForegroundEstimate now exposes rgb (F), background (B) and a correction map;
  debug outputs become foreground_rgb / foreground_background / foreground_correction.

Controlled comparison (same ViTMatte alpha, only foreground method changed),
edge_green_excess:
  TestImage  p95 0.176->0.031, mean 0.061->0.005
  TestImage2 p95 0.165->0.004, mean 0.052->0.002

Adds pymatting (pulls in numba) to requirements.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

README.md

@@ -84,10 +84,9 @@ For each processed image, the CLI writes an RGBA PNG and optional debug files:
 bg_confidence.png
 trimap.png
 alpha.png
-foreground_mask.png
-foreground_unmix_mask.png
-foreground_local_mask.png
 foreground_rgb.png
+foreground_background.png
+foreground_correction.png
 despill_mask.png
 preview_black.png
 preview_white.png
@@ -125,5 +124,7 @@ D:\MiniConda\envs\lightML\python.exe scripts\smoke_samples.py `
 
 - `Samples/` and `Outputs/` are ignored by Git.
 - `ViTMatte` model weights are loaded from Hugging Face on first use.
-- Foreground color estimation runs before final de-spill and writes debug masks
-  for unmixing and local foreground propagation.
+- Foreground color estimation uses pymatting's `estimate_foreground_ml` to
+  propagate clean foreground colour into semi-transparent edges before final
+  de-spill, writing the estimated foreground, background, and a correction map.
+  Set `foreground.method: unmix` to fall back to the legacy heuristic.

bgfilter/foreground.py

@@ -11,10 +11,72 @@ from .settings import ForegroundSettings
 
 @dataclass(frozen=True)
 class ForegroundEstimate:
-    rgb: np.ndarray
-    mask: np.ndarray
-    unmix_mask: np.ndarray
-    local_mask: np.ndarray
+    rgb: np.ndarray          # estimated foreground colour F (uint8, HxWx3)
+    background: np.ndarray   # estimated background colour B (uint8, HxWx3)
+    correction: np.ndarray   # per-pixel |F - input| magnitude, 0..1 (float32, HxW)
+
+
+def estimate_foreground_rgb(
+    rgb: np.ndarray,
+    alpha: np.ndarray,
+    bg_confidence: np.ndarray,
+    model: BackgroundModel,
+    settings: ForegroundSettings,
+) -> ForegroundEstimate:
+    if not settings.enabled:
+        bg = np.broadcast_to(np.asarray(model.rgb_center, dtype=np.float32), rgb.shape)
+        return ForegroundEstimate(
+            rgb=rgb.copy(),
+            background=np.clip(bg * 255.0, 0, 255).astype(np.uint8),
+            correction=np.zeros(alpha.shape, dtype=np.float32),
+        )
+
+    method = settings.method
+    if method == "ml":
+        try:
+            return _estimate_ml(rgb, alpha, settings)
+        except RuntimeError:
+            if not settings.fallback_to_unmix:
+                raise
+            method = "unmix"
+    if method == "unmix":
+        return _estimate_unmix(rgb, alpha, bg_confidence, model, settings)
+    raise RuntimeError(
+        f"Unsupported foreground method '{settings.method}'. Use 'ml' or 'unmix'."
+    )
+
+
+def _estimate_ml(
+    rgb: np.ndarray, alpha: np.ndarray, settings: ForegroundSettings
+) -> ForegroundEstimate:
+    """Multi-level closed-form foreground/background estimation (pymatting).
+
+    Solves image = alpha * F + (1 - alpha) * B with spatial smoothness, so the
+    foreground colour is *propagated* into semi-transparent edges from reliable
+    high-alpha neighbours instead of being divided out pixel-wise. That removes
+    green spill on thin hair where a direct unmix (divide by small alpha) is
+    too noisy to trust.
+    """
+    try:
+        from pymatting import estimate_foreground_ml
+    except ModuleNotFoundError as exc:
+        raise RuntimeError(
+            "Missing dependency 'pymatting'. Install it with: pip install pymatting"
+        ) from exc
+
+    image = rgb.astype(np.float64) / 255.0
+    a = np.clip(alpha.astype(np.float64), 0.0, 1.0)
+    foreground, background = estimate_foreground_ml(
+        image, a, regularization=settings.ml_regularization, return_background=True
+    )
+    foreground = np.clip(foreground, 0.0, 1.0)
+    background = np.clip(background, 0.0, 1.0)
+    correction = np.linalg.norm(foreground - image, axis=2) / np.sqrt(3.0)
+    return ForegroundEstimate(
+        rgb=np.clip(foreground * 255.0, 0, 255).astype(np.uint8),
+        background=np.clip(background * 255.0, 0, 255).astype(np.uint8),
+        correction=correction.astype(np.float32),
+    )
 
 
 def _weighted_blur(values: np.ndarray, weights: np.ndarray, radius: int) -> np.ndarray:
@@ -29,18 +91,15 @@ def _weighted_blur(values: np.ndarray, weights: np.ndarray, radius: int) -> np.n
     return numerator / np.maximum(denominator[..., None], 1e-6)
 
 
-def estimate_foreground_rgb(
+def _estimate_unmix(
     rgb: np.ndarray,
     alpha: np.ndarray,
     bg_confidence: np.ndarray,
     model: BackgroundModel,
     settings: ForegroundSettings,
 ) -> ForegroundEstimate:
+    """Legacy heuristic: edge-band alpha unmix plus local foreground blur."""
     rgb_f = rgb.astype(np.float32) / 255.0
-    zero = np.zeros(alpha.shape, dtype=np.float32)
-    if not settings.enabled:
-        return ForegroundEstimate(rgb=rgb.copy(), mask=zero, unmix_mask=zero, local_mask=zero)
-
     r = rgb_f[..., 0]
     g = rgb_f[..., 1]
     b = rgb_f[..., 2]
@@ -72,10 +131,10 @@ def estimate_foreground_rgb(
     local_mask = base_mask * settings.local_strength * low_alpha_weight
     out = out * (1.0 - local_mask[..., None]) + local_fg * local_mask[..., None]
 
-    combined_mask = np.clip(unmix_mask + local_mask, 0.0, 1.0)
+    background = np.broadcast_to(bg, rgb_f.shape)
+    correction = np.linalg.norm(out - rgb_f, axis=2) / np.sqrt(3.0)
     return ForegroundEstimate(
         rgb=np.clip(out * 255.0, 0, 255).astype(np.uint8),
-        mask=combined_mask.astype(np.float32),
-        unmix_mask=unmix_mask.astype(np.float32),
-        local_mask=local_mask.astype(np.float32),
+        background=np.clip(background * 255.0, 0, 255).astype(np.uint8),
+        correction=correction.astype(np.float32),
     )

bgfilter/pipeline.py

@@ -95,10 +95,9 @@ def _run_image(
         save_gray(debug / "bg_confidence.png", bg_confidence)
         save_gray(debug / "trimap.png", trimap)
         save_gray(debug / "alpha.png", alpha)
-        save_gray(debug / "foreground_mask.png", foreground.mask)
-        save_gray(debug / "foreground_unmix_mask.png", foreground.unmix_mask)
-        save_gray(debug / "foreground_local_mask.png", foreground.local_mask)
         save_rgb(debug / "foreground_rgb.png", foreground.rgb)
+        save_rgb(debug / "foreground_background.png", foreground.background)
+        save_gray(debug / "foreground_correction.png", foreground.correction)
         save_gray(debug / "despill_mask.png", despill_mask)
         save_previews(debug, corrected_rgb, alpha)
         tiles = {

bgfilter/settings.py

@@ -38,6 +38,10 @@ class AlphaPostSettings:
 @dataclass(frozen=True)
 class ForegroundSettings:
     enabled: bool = True
+    method: str = "ml"  # "ml" (pymatting) or "unmix" (legacy heuristic)
+    fallback_to_unmix: bool = True
+    ml_regularization: float = 1e-5
+    # Legacy "unmix" method parameters (used only when method == "unmix").
     edge_low: float = 0.005
     edge_high: float = 0.995
     min_unmix_alpha: float = 0.08

configs/default.yaml

@@ -12,6 +12,10 @@ trimap:
 
 foreground:
   enabled: true
+  method: ml
+  fallback_to_unmix: true
+  ml_regularization: 0.00001
+  # Legacy "unmix" method params (used only when method: unmix).
   edge_low: 0.005
   edge_high: 0.995
   min_unmix_alpha: 0.08

docs/green_screen_matting_workflow.md

@@ -372,10 +372,9 @@ result_rgba.png
 debug/bg_confidence.png
 debug/trimap.png
 debug/alpha.png
-debug/foreground_mask.png
-debug/foreground_unmix_mask.png
-debug/foreground_local_mask.png
 debug/foreground_rgb.png
+debug/foreground_background.png
+debug/foreground_correction.png
 debug/despill_mask.png
 debug/preview_black.png
 debug/preview_white.png

requirements.txt

@@ -3,6 +3,7 @@ pillow
 numpy
 scipy
 scikit-image
+pymatting
 torch
 transformers
 accelerate

scripts/smoke_samples.py

@@ -19,10 +19,9 @@ REQUIRED_DEBUG_FILES = {
     "trimap.png",
     "alpha.png",
     "despill_mask.png",
-    "foreground_mask.png",
-    "foreground_unmix_mask.png",
-    "foreground_local_mask.png",
     "foreground_rgb.png",
+    "foreground_background.png",
+    "foreground_correction.png",
     "preview_black.png",
     "preview_white.png",
     "preview_gray.png",