From e234a1a13d553f7a3153fd996629aa3565c6d67b Mon Sep 17 00:00:00 2001 From: lingh Date: Tue, 7 Jul 2026 19:49:53 +0800 Subject: [PATCH] Backport server-edition's pipeline structure to master Adopt the service branch's refactors verbatim so the two branches stop drifting on these files: - pipeline.py: split _run_image into _process_rgb + MattingResult (an in-memory API, exposed as MattingPipeline.run_rgb) and allow injecting pre-loaded models into MattingPipeline. Disk I/O and debug dumps stay in _run_image; the algorithm is line-for-line the same. - foreground.py: estimate_foreground_rgb returns the foreground array directly; the ForegroundEstimate dataclass carried background/ correction fields nothing consumed. Verified: default TestImage3 run byte-identical to pre-backport master, all debug artifacts written, run_rgb output matches run_image bytes under save_rgba's quantization; both files now diff-clean against server-edition. Co-Authored-By: Claude Opus 4.8 --- bgfilter/foreground.py | 44 +++------------- bgfilter/pipeline.py | 113 ++++++++++++++++++++++++++++------------- 2 files changed, 87 insertions(+), 70 deletions(-) diff --git a/bgfilter/foreground.py b/bgfilter/foreground.py index a26f2b3..b898761 100644 --- a/bgfilter/foreground.py +++ b/bgfilter/foreground.py @@ -1,7 +1,5 @@ from __future__ import annotations -from dataclasses import dataclass - import numpy as np from .chroma import BackgroundModel @@ -9,27 +7,16 @@ from .deps import require_cv2 from .settings import ForegroundSettings -@dataclass(frozen=True) -class ForegroundEstimate: - 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: +) -> np.ndarray: + """Estimated foreground colour F (uint8 HxWx3) for compositing over alpha.""" 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), - ) + return rgb.copy() method = settings.method if method == "ml": @@ -48,7 +35,7 @@ def estimate_foreground_rgb( def _estimate_ml( rgb: np.ndarray, alpha: np.ndarray, settings: ForegroundSettings -) -> ForegroundEstimate: +) -> np.ndarray: """Multi-level closed-form foreground/background estimation (pymatting). Solves image = alpha * F + (1 - alpha) * B with spatial smoothness, so the @@ -66,17 +53,8 @@ def _estimate_ml( 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), - ) + foreground = estimate_foreground_ml(image, a, regularization=settings.ml_regularization) + return np.clip(foreground * 255.0, 0, 255).astype(np.uint8) def _weighted_blur(values: np.ndarray, weights: np.ndarray, radius: int) -> np.ndarray: @@ -97,7 +75,7 @@ def _estimate_unmix( bg_confidence: np.ndarray, model: BackgroundModel, settings: ForegroundSettings, -) -> ForegroundEstimate: +) -> np.ndarray: """Legacy heuristic: edge-band alpha unmix plus local foreground blur.""" rgb_f = rgb.astype(np.float32) / 255.0 r = rgb_f[..., 0] @@ -131,10 +109,4 @@ def _estimate_unmix( local_mask = base_mask * settings.local_strength * low_alpha_weight out = out * (1.0 - local_mask[..., None]) + local_fg * local_mask[..., None] - 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), - background=np.clip(background * 255.0, 0, 255).astype(np.uint8), - correction=correction.astype(np.float32), - ) + return np.clip(out * 255.0, 0, 255).astype(np.uint8) diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 357dafe..f1906aa 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from dataclasses import asdict +from dataclasses import asdict, dataclass from pathlib import Path import numpy as np @@ -32,12 +32,41 @@ from .trimap import ( from .vitmatte_infer import ViTMatteRunner +@dataclass +class MattingResult: + """In-memory result of one matting run. + + ``rgb`` is the despilled foreground (uint8 HxWx3) and ``alpha`` is float32 HxW in + 0..1; together they compose the output RGBA. The remaining arrays are the + intermediate maps used for debug previews. + """ + + rgb: np.ndarray + alpha: np.ndarray + metadata: dict + bg_confidence: np.ndarray + trimap: np.ndarray + seg_mask: np.ndarray | None + cross_check_alpha: np.ndarray | None + color_mask: np.ndarray + input_rgb: np.ndarray + + class MattingPipeline: - def __init__(self, settings: PipelineSettings): + def __init__( + self, + settings: PipelineSettings, + *, + runner: ViTMatteRunner | None = None, + segmenter=None, + cross_checker=None, + ): self.settings = settings - self._runner: ViTMatteRunner | None = None - self._segmenter = None - self._cross_checker = None + # Pre-loaded models can be injected so a long-lived service shares them + # across requests instead of reloading per call; None = lazy-load on first use. + self._runner: ViTMatteRunner | None = runner + self._segmenter = segmenter + self._cross_checker = cross_checker def _segment(self, rgb: np.ndarray) -> np.ndarray: if self._segmenter is None: @@ -79,6 +108,10 @@ class MattingPipeline: raise return trimap_to_alpha_seed(trimap, bg_confidence), "chroma_fallback" + def run_rgb(self, rgb: np.ndarray) -> MattingResult: + """Run the full pipeline on an in-memory RGB array (no disk I/O).""" + return _process_rgb(rgb, self) + def run_image( self, input_path: str | Path, @@ -97,14 +130,8 @@ def run_image( return MattingPipeline(settings).run_image(input_path, output_path, debug_dir) -def _run_image( - input_path: str | Path, - output_path: str | Path, - debug_dir: str | Path | None, - pipeline: MattingPipeline, -) -> dict: +def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: settings = pipeline.settings - rgb = load_rgb(input_path) screen_color = parse_hex_color(settings.screen_color) if settings.screen_color else None bg_confidence, model = compute_bg_confidence( rgb, settings=settings.chroma, screen_color=screen_color @@ -168,41 +195,59 @@ def _run_image( foreground = estimate_foreground_rgb( rgb, alpha, bg_confidence, model, settings.foreground ) - corrected_rgb, color_mask = despill( - foreground.rgb, alpha, model, settings.despill - ) - save_rgba(output_path, corrected_rgb, alpha) + corrected_rgb, color_mask = despill(foreground, alpha, model, settings.despill) - result = { - "input": str(input_path), - "output": str(output_path), + metadata = { "alpha_source": alpha_source, "background_model": model.to_dict(), "trimap": trimap_stats, "settings": asdict(settings), } + return MattingResult( + rgb=corrected_rgb, + alpha=alpha, + metadata=metadata, + bg_confidence=bg_confidence, + trimap=trimap, + seg_mask=seg_mask, + cross_check_alpha=second_alpha, + color_mask=color_mask, + input_rgb=rgb, + ) + + +def _run_image( + input_path: str | Path, + output_path: str | Path, + debug_dir: str | Path | None, + pipeline: MattingPipeline, +) -> dict: + res = _process_rgb(load_rgb(input_path), pipeline) + save_rgba(output_path, res.rgb, res.alpha) + + result = {"input": str(input_path), "output": str(output_path), **res.metadata} if debug_dir is not None: debug = Path(debug_dir) debug.mkdir(parents=True, exist_ok=True) - save_gray(debug / "bg_confidence.png", bg_confidence) - save_gray(debug / "trimap.png", trimap) - if seg_mask is not None: - save_gray(debug / "seg_mask.png", seg_mask) - if second_alpha is not None: - save_gray(debug / "cross_check_alpha.png", second_alpha) - save_gray(debug / "alpha.png", alpha) - save_rgb(debug / "foreground_rgb.png", corrected_rgb) - save_gray(debug / "color_mask.png", color_mask) - save_previews(debug, corrected_rgb, alpha) + save_gray(debug / "bg_confidence.png", res.bg_confidence) + save_gray(debug / "trimap.png", res.trimap) + if res.seg_mask is not None: + save_gray(debug / "seg_mask.png", res.seg_mask) + if res.cross_check_alpha is not None: + save_gray(debug / "cross_check_alpha.png", res.cross_check_alpha) + save_gray(debug / "alpha.png", res.alpha) + save_rgb(debug / "foreground_rgb.png", res.rgb) + save_gray(debug / "color_mask.png", res.color_mask) + save_previews(debug, res.rgb, res.alpha) tiles = { - "input": rgb, - "bg confidence": np.repeat((bg_confidence[..., None] * 255).astype(np.uint8), 3, axis=2), - "trimap": np.repeat(trimap[..., None], 3, axis=2), - "alpha": np.repeat((alpha[..., None] * 255).astype(np.uint8), 3, axis=2), + "input": res.input_rgb, + "bg confidence": np.repeat((res.bg_confidence[..., None] * 255).astype(np.uint8), 3, axis=2), + "trimap": np.repeat(res.trimap[..., None], 3, axis=2), + "alpha": np.repeat((res.alpha[..., None] * 255).astype(np.uint8), 3, axis=2), } for name, color in BACKGROUND_COLORS.items(): - tiles[f"preview {name}"] = composite(corrected_rgb, alpha, color) + tiles[f"preview {name}"] = composite(res.rgb, res.alpha, color) make_qa_grid(debug, tiles) write_text(debug / "metadata.json", json.dumps(result, indent=2, ensure_ascii=False))