From 923dbe9f09b516beb47715485b2bcf59ba3519ab Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 30 Jun 2026 14:42:08 +0800 Subject: [PATCH] Add sample smoke check and refine foreground despill --- README.md | 13 +++- bgfilter/despill.py | 24 ++++++ bgfilter/settings.py | 4 + configs/default.yaml | 4 + docs/green_screen_matting_workflow.md | 1 + scripts/smoke_samples.py | 101 ++++++++++++++++++++++++++ 6 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 scripts/smoke_samples.py diff --git a/README.md b/README.md index 66f5f43..c28f765 100644 --- a/README.md +++ b/README.md @@ -94,7 +94,18 @@ edge pixels. ```powershell D:\MiniConda\envs\lightML\python.exe -m bgfilter.quality_cli ` Outputs\TestImage_rgba.png ` - --max-edge-green-excess-p95 0.55 + --max-edge-green-excess-p95 0.50 +``` + +Run the bundled sample smoke check: + +```powershell +D:\MiniConda\envs\lightML\python.exe scripts\smoke_samples.py ` + --samples-dir Samples ` + --output-dir Outputs\smoke_samples ` + --device cpu ` + --fallback-to-chroma-alpha ` + --max-edge-green-excess-p95 0.50 ``` ## Current Notes diff --git a/bgfilter/despill.py b/bgfilter/despill.py index 7d65f38..3a032f5 100644 --- a/bgfilter/despill.py +++ b/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) diff --git a/bgfilter/settings.py b/bgfilter/settings.py index d90dd00..3cd3d86 100644 --- a/bgfilter/settings.py +++ b/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) diff --git a/configs/default.yaml b/configs/default.yaml index 98d6a88..4bd1c1d 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -21,3 +21,7 @@ despill: alpha_weight_floor: 0.35 unmix_strength: 0.55 unmix_alpha_min: 0.08 + local_fg_strength: 0.35 + local_fg_blur_radius: 9 + local_fg_alpha_threshold: 0.92 + local_fg_bg_confidence_max: 0.20 diff --git a/docs/green_screen_matting_workflow.md b/docs/green_screen_matting_workflow.md index a92fcf3..4c10610 100644 --- a/docs/green_screen_matting_workflow.md +++ b/docs/green_screen_matting_workflow.md @@ -424,6 +424,7 @@ debug/qa_grid.png 5. 发丝区域不是硬切边。 6. 镂空洞露出绿幕时可透明。 7. 边缘没有明显绿色污染。 +8. `scripts/smoke_samples.py` 能处理样例图并验证 RGBA、debug 产物和边缘绿色污染指标。 ``` ## 19. 推荐落地顺序 diff --git a/scripts/smoke_samples.py b/scripts/smoke_samples.py new file mode 100644 index 0000000..0061e3b --- /dev/null +++ b/scripts/smoke_samples.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from bgfilter.pipeline import MattingPipeline +from bgfilter.quality import assert_quality, measure_rgba +from bgfilter.settings import ( + AlphaPostSettings, + ChromaSettings, + DespillSettings, + ModelSettings, + PipelineSettings, + TrimapSettings, +) + + +REQUIRED_DEBUG_FILES = { + "bg_confidence.png", + "trimap.png", + "alpha.png", + "despill_mask.png", + "preview_black.png", + "preview_white.png", + "preview_gray.png", + "preview_red.png", + "preview_blue.png", + "qa_grid.png", + "metadata.json", +} + + +def build_settings(args: argparse.Namespace) -> PipelineSettings: + return PipelineSettings( + chroma=ChromaSettings(), + trimap=TrimapSettings(), + alpha_post=AlphaPostSettings(), + despill=DespillSettings(), + model=ModelSettings( + model_name=args.model_name, + device=args.device, + matting_method=args.matting_method, + fallback_to_chroma_alpha=args.fallback_to_chroma_alpha, + ), + ) + + +def check_debug_files(debug_dir: Path) -> list[str]: + missing = [] + for filename in sorted(REQUIRED_DEBUG_FILES): + path = debug_dir / filename + if not path.exists() or path.stat().st_size <= 0: + missing.append(filename) + return missing + + +def main() -> int: + parser = argparse.ArgumentParser(description="Run sample-image smoke checks.") + parser.add_argument("--samples-dir", type=Path, default=Path("Samples")) + parser.add_argument("--output-dir", type=Path, default=Path("Outputs/smoke_samples")) + parser.add_argument("--device", default="cpu") + parser.add_argument("--matting-method", default="vitmatte", choices=["vitmatte", "chroma"]) + parser.add_argument("--model-name", default="hustvl/vitmatte-small-composition-1k") + parser.add_argument("--fallback-to-chroma-alpha", action="store_true") + parser.add_argument("--max-edge-green-excess-p95", type=float, default=0.50) + args = parser.parse_args() + + images = sorted(args.samples_dir.glob("*.png")) + if not images: + raise SystemExit(f"No PNG samples found in {args.samples_dir}") + + pipeline = MattingPipeline(build_settings(args)) + failed = False + for image_path in images: + output_path = args.output_dir / f"{image_path.stem}_rgba.png" + debug_dir = args.output_dir / f"{image_path.stem}_debug" + result = pipeline.run_image(image_path, output_path, debug_dir) + metrics = measure_rgba(output_path) + failures = assert_quality(metrics, args.max_edge_green_excess_p95) + missing_debug = check_debug_files(debug_dir) + if missing_debug: + failures.append(f"missing debug files: {', '.join(missing_debug)}") + + print( + f"{image_path.name}: alpha_source={result['alpha_source']} " + f"edge_green_p95={metrics.edge_green_excess_p95:.4f}" + ) + for failure in failures: + print(f" FAIL {failure}") + failed = failed or bool(failures) + + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main())