diff --git a/README.md b/README.md index fd2a20e..75c8bea 100644 --- a/README.md +++ b/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. diff --git a/bgfilter/foreground.py b/bgfilter/foreground.py index bca3f30..ff79dd3 100644 --- a/bgfilter/foreground.py +++ b/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), ) diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index bc385be..005903a 100644 --- a/bgfilter/pipeline.py +++ b/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 = { diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 84541c0..d3601ba 100644 --- a/bgfilter/settings.py +++ b/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 diff --git a/configs/default.yaml b/configs/default.yaml index 279036d..c187e37 100644 --- a/configs/default.yaml +++ b/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 diff --git a/docs/green_screen_matting_workflow.md b/docs/green_screen_matting_workflow.md index 6b4958f..3fc6fb4 100644 --- a/docs/green_screen_matting_workflow.md +++ b/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 diff --git a/requirements.txt b/requirements.txt index 35d5373..c48f6f6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ pillow numpy scipy scikit-image +pymatting torch transformers accelerate diff --git a/scripts/smoke_samples.py b/scripts/smoke_samples.py index 96ca103..f864ec5 100644 --- a/scripts/smoke_samples.py +++ b/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",