Add foreground color estimation pass

bgfilter/pipeline.py

@@ -9,7 +9,8 @@ import numpy as np
 from .alpha_post import clean_alpha, enforce_trimap
 from .chroma import compute_bg_confidence
 from .despill import despill_green
-from .io import load_rgb, save_gray, save_rgba, write_text
+from .foreground import estimate_foreground_rgb
+from .io import load_rgb, save_gray, save_rgb, save_rgba, write_text
 from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews
 from .settings import PipelineSettings
 from .trimap import generate_trimap, trimap_to_alpha_seed
@@ -71,8 +72,11 @@ def _run_image(
 
     alpha = enforce_trimap(alpha, trimap)
     alpha = clean_alpha(alpha, trimap, settings.alpha_post)
+    foreground = estimate_foreground_rgb(
+        rgb, alpha, bg_confidence, model, settings.foreground
+    )
     corrected_rgb, despill_mask = despill_green(
-        rgb, alpha, bg_confidence, model, settings.despill
+        foreground.rgb, alpha, bg_confidence, model, settings.despill
     )
     save_rgba(output_path, corrected_rgb, alpha)
 
@@ -91,6 +95,10 @@ 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_gray(debug / "despill_mask.png", despill_mask)
         save_previews(debug, corrected_rgb, alpha)
         tiles = {