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 <noreply@anthropic.com>
This commit is contained in:
+8
-36
@@ -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)
|
||||
|
||||
+79
-34
@@ -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))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user