From 464a60c733c0af3ee023f53bcd910be5d019bc3a Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 30 Jun 2026 14:29:19 +0800 Subject: [PATCH] Improve despill and add quality checks --- README.md | 16 +++++++-- bgfilter/despill.py | 45 ++++++++++++++++++++--- bgfilter/quality.py | 80 +++++++++++++++++++++++++++++++++++++++++ bgfilter/quality_cli.py | 38 ++++++++++++++++++++ bgfilter/settings.py | 13 ++++--- configs/default.yaml | 12 +++++-- 6 files changed, 190 insertions(+), 14 deletions(-) create mode 100644 bgfilter/quality.py create mode 100644 bgfilter/quality_cli.py diff --git a/README.md b/README.md index b9c7c90..66f5f43 100644 --- a/README.md +++ b/README.md @@ -86,9 +86,21 @@ qa_grid.png metadata.json ``` +## Quality Check + +The quality checker measures alpha validity and green spill on semi-transparent +edge pixels. + +```powershell +D:\MiniConda\envs\lightML\python.exe -m bgfilter.quality_cli ` + Outputs\TestImage_rgba.png ` + --max-edge-green-excess-p95 0.55 +``` + ## Current Notes - `Samples/` and `Outputs/` are ignored by Git. - `ViTMatte` model weights are loaded from Hugging Face on first use. -- Current despill is conservative and edge-local. Strong green spill around hair may - still need parameter tuning or a stronger foreground color estimation pass. +- Current despill is edge-local and includes a lightweight foreground unmixing pass. + Strong green spill around hair may still need more advanced foreground color + estimation. diff --git a/bgfilter/despill.py b/bgfilter/despill.py index c38f875..7d65f38 100644 --- a/bgfilter/despill.py +++ b/bgfilter/despill.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np from .chroma import BackgroundModel +from .deps import require_cv2 from .settings import DespillSettings @@ -17,23 +18,57 @@ def despill_green( if not settings.enabled: return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32) + cv2 = require_cv2() edge = (alpha > settings.edge_low) & (alpha < settings.edge_high) + if settings.edge_expand_radius > 0: + radius = settings.edge_expand_radius + kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1) + ) + edge = cv2.dilate(edge.astype(np.uint8), kernel).astype(bool) + r = rgb_f[..., 0] g = rgb_f[..., 1] b = rgb_f[..., 2] neutral_green = np.maximum(r, b) + settings.green_excess_margin excess = np.maximum(g - neutral_green, 0.0) - edge_weight = np.clip((1.0 - np.abs(alpha - 0.5) * 2.0), 0.0, 1.0) - bg_weight = np.clip(bg_confidence, 0.0, 1.0) - mask = edge.astype(np.float32) * edge_weight * np.maximum(bg_weight, 0.25) + alpha_weight = np.maximum( + np.clip((1.0 - np.abs(alpha - 0.5) * 2.0), 0.0, 1.0), + settings.alpha_weight_floor, + ) + bg_weight = ( + settings.bg_confidence_weight * np.clip(bg_confidence, 0.0, 1.0) + + (1.0 - settings.bg_confidence_weight) + ) + green_weight = np.clip(excess / max(settings.green_excess_margin * 8.0, 1e-6), 0.0, 1.0) + mask = edge.astype(np.float32) * alpha_weight * bg_weight * green_weight out = rgb_f.copy() - out[..., 1] = g - excess * mask * settings.strength + target_green = np.minimum(g, neutral_green) + out[..., 1] = g * (1.0 - mask * settings.strength) + target_green * ( + mask * settings.strength + ) - # Very light compensation toward the non-green channels to avoid gray fringes. + # Light foreground color approximation: green spill usually removed too much + # luminance from thin hair edges, so redistribute part of the removed green + # into red/blue without changing fully opaque pixels. bg_green = float(model.rgb_center[1]) compensation = excess * mask * settings.strength * min(0.25, bg_green * 0.15) out[..., 0] = np.clip(out[..., 0] + compensation * 0.5, 0.0, 1.0) out[..., 2] = np.clip(out[..., 2] + compensation * 0.5, 0.0, 1.0) + # Approximate foreground recovery for semi-transparent contaminated pixels. + # C = alpha * F + (1 - alpha) * B, so F = (C - (1 - alpha) * B) / alpha. + # Blend it only where green excess is present; this is intentionally local + # because incorrect unmixing on opaque clothing would damage normal colors. + bg = np.asarray(model.rgb_center, dtype=np.float32) + a = np.clip(alpha[..., None], settings.unmix_alpha_min, 1.0) + recovered = np.clip((rgb_f - (1.0 - alpha[..., None]) * bg) / a, 0.0, 1.0) + unmix_mask = ( + mask + * settings.unmix_strength + * np.clip((1.0 - alpha) / max(1.0 - settings.unmix_alpha_min, 1e-6), 0.0, 1.0) + ) + out = out * (1.0 - unmix_mask[..., None]) + recovered * unmix_mask[..., None] + return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32) diff --git a/bgfilter/quality.py b/bgfilter/quality.py new file mode 100644 index 0000000..48b6151 --- /dev/null +++ b/bgfilter/quality.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from pathlib import Path + +import numpy as np +from PIL import Image + + +@dataclass(frozen=True) +class QualityMetrics: + path: str + width: int + height: int + alpha_min: int + alpha_max: int + alpha_mean: float + nonzero_alpha_pixels: int + edge_pixels: int + edge_green_excess_mean: float + edge_green_excess_p95: float + edge_green_excess_max: float + + def to_dict(self) -> dict: + return asdict(self) + + +def measure_rgba(path: str | Path) -> QualityMetrics: + image = Image.open(path).convert("RGBA") + arr = np.asarray(image).astype(np.float32) / 255.0 + rgb = arr[..., :3] + alpha = arr[..., 3] + edge = (alpha > 0.02) & (alpha < 0.98) + green_excess = np.maximum(rgb[..., 1] - np.maximum(rgb[..., 0], rgb[..., 2]), 0.0) + edge_values = green_excess[edge] + if edge_values.size: + edge_mean = float(edge_values.mean()) + edge_p95 = float(np.percentile(edge_values, 95)) + edge_max = float(edge_values.max()) + else: + edge_mean = edge_p95 = edge_max = 0.0 + + alpha_u8 = (alpha * 255.0).round().astype(np.uint8) + return QualityMetrics( + path=str(path), + width=image.width, + height=image.height, + alpha_min=int(alpha_u8.min()), + alpha_max=int(alpha_u8.max()), + alpha_mean=float(alpha_u8.mean()), + nonzero_alpha_pixels=int((alpha_u8 > 0).sum()), + edge_pixels=int(edge.sum()), + edge_green_excess_mean=edge_mean, + edge_green_excess_p95=edge_p95, + edge_green_excess_max=edge_max, + ) + + +def assert_quality( + metrics: QualityMetrics, + max_edge_green_excess_p95: float | None = None, + require_alpha_range: bool = True, +) -> list[str]: + failures: list[str] = [] + if require_alpha_range: + if metrics.alpha_min != 0: + failures.append(f"alpha_min is {metrics.alpha_min}, expected 0") + if metrics.alpha_max != 255: + failures.append(f"alpha_max is {metrics.alpha_max}, expected 255") + if metrics.nonzero_alpha_pixels <= 0: + failures.append("nonzero alpha pixel count is 0") + if ( + max_edge_green_excess_p95 is not None + and metrics.edge_green_excess_p95 > max_edge_green_excess_p95 + ): + failures.append( + "edge_green_excess_p95 " + f"{metrics.edge_green_excess_p95:.4f} exceeds {max_edge_green_excess_p95:.4f}" + ) + return failures diff --git a/bgfilter/quality_cli.py b/bgfilter/quality_cli.py new file mode 100644 index 0000000..9864b03 --- /dev/null +++ b/bgfilter/quality_cli.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import typer +from rich import print + +from .quality import assert_quality, measure_rgba + +app = typer.Typer(help="Quality checks for generated RGBA PNG files.") + + +@app.command() +def main( + image: Path = typer.Argument(..., exists=True, file_okay=True, dir_okay=False), + max_edge_green_excess_p95: float | None = typer.Option( + None, "--max-edge-green-excess-p95", min=0.0, max=1.0 + ), + json_output: bool = typer.Option(False, "--json/--no-json"), +) -> None: + metrics = measure_rgba(image) + failures = assert_quality(metrics, max_edge_green_excess_p95) + + if json_output: + print(json.dumps(metrics.to_dict(), indent=2)) + else: + print(metrics.to_dict()) + + if failures: + for failure in failures: + print(f"[red]fail:[/red] {failure}") + raise typer.Exit(code=1) + print("[green]quality check passed[/green]") + + +if __name__ == "__main__": + app() diff --git a/bgfilter/settings.py b/bgfilter/settings.py index a31435b..d90dd00 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -38,10 +38,15 @@ class AlphaPostSettings: @dataclass(frozen=True) class DespillSettings: enabled: bool = True - edge_low: float = 0.02 - edge_high: float = 0.98 - strength: float = 0.75 - green_excess_margin: float = 0.03 + edge_low: float = 0.005 + edge_high: float = 0.995 + strength: float = 0.92 + green_excess_margin: float = 0.015 + edge_expand_radius: int = 2 + bg_confidence_weight: float = 0.65 + alpha_weight_floor: float = 0.35 + unmix_strength: float = 0.55 + unmix_alpha_min: float = 0.08 @dataclass(frozen=True) diff --git a/configs/default.yaml b/configs/default.yaml index a953a45..98d6a88 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -12,6 +12,12 @@ trimap: despill: enabled: true - edge_low: 0.02 - edge_high: 0.98 - strength: 0.75 + edge_low: 0.005 + edge_high: 0.995 + strength: 0.92 + green_excess_margin: 0.015 + edge_expand_radius: 2 + bg_confidence_weight: 0.65 + alpha_weight_floor: 0.35 + unmix_strength: 0.55 + unmix_alpha_min: 0.08