Implement green screen matting pipeline
This commit is contained in:
Codex
@@ -1,2 +1,4 @@
|
||||
Samples/
|
||||
Outputs/
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# BgFilter
|
||||
|
||||
Offline green screen character matting for AI-generated character images.
|
||||
|
||||
The first implementation follows the workflow in
|
||||
[`docs/green_screen_matting_workflow.md`](docs/green_screen_matting_workflow.md):
|
||||
|
||||
```text
|
||||
RGB input -> chroma confidence -> trimap -> ViTMatte -> alpha cleanup -> despill -> RGBA PNG -> QA previews
|
||||
```
|
||||
|
||||
## Environment
|
||||
|
||||
Use the conda environment `lightML`.
|
||||
|
||||
```powershell
|
||||
conda activate lightML
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
If the shell is not activated, call the environment Python directly:
|
||||
|
||||
```powershell
|
||||
D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli --help
|
||||
```
|
||||
|
||||
## Single Image
|
||||
|
||||
```powershell
|
||||
D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli `
|
||||
--input Samples\TestImage.png `
|
||||
--output Outputs\TestImage_rgba.png `
|
||||
--debug-dir Outputs\TestImage_debug `
|
||||
--device cuda
|
||||
```
|
||||
|
||||
Use CPU for validation when CUDA is unavailable:
|
||||
|
||||
```powershell
|
||||
D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli `
|
||||
--input Samples\TestImage.png `
|
||||
--output Outputs\TestImage_rgba.png `
|
||||
--debug-dir Outputs\TestImage_debug `
|
||||
--device cpu
|
||||
```
|
||||
|
||||
## Batch
|
||||
|
||||
```powershell
|
||||
D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli `
|
||||
--input-dir Samples `
|
||||
--output-dir Outputs `
|
||||
--debug-dir Outputs\debug `
|
||||
--device cuda
|
||||
```
|
||||
|
||||
## Chroma-Only Debug Mode
|
||||
|
||||
This mode skips ViTMatte and uses chroma confidence as an alpha seed. It is useful
|
||||
for fast debugging of chroma confidence, trimap, despill, and QA outputs.
|
||||
|
||||
```powershell
|
||||
D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli `
|
||||
--input-dir Samples `
|
||||
--output-dir Outputs\chroma `
|
||||
--debug-dir Outputs\chroma_debug `
|
||||
--matting-method chroma `
|
||||
--device cpu
|
||||
```
|
||||
|
||||
## Outputs
|
||||
|
||||
For each processed image, the CLI writes an RGBA PNG and optional debug files:
|
||||
|
||||
```text
|
||||
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
|
||||
```
|
||||
|
||||
## 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.
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Green screen character matting pipeline."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,55 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .deps import require_cv2
|
||||
from .settings import AlphaPostSettings
|
||||
|
||||
|
||||
def enforce_trimap(alpha: np.ndarray, trimap: np.ndarray) -> np.ndarray:
|
||||
out = np.clip(alpha.astype(np.float32), 0.0, 1.0)
|
||||
out[trimap == 0] = 0.0
|
||||
out[trimap == 255] = 1.0
|
||||
return out
|
||||
|
||||
|
||||
def _remove_small_components(mask: np.ndarray, min_area: int) -> np.ndarray:
|
||||
cv2 = require_cv2()
|
||||
count, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8)
|
||||
out = np.zeros_like(mask, dtype=bool)
|
||||
for label in range(1, count):
|
||||
if stats[label, cv2.CC_STAT_AREA] >= min_area:
|
||||
out[labels == label] = True
|
||||
return out
|
||||
|
||||
|
||||
def clean_alpha(
|
||||
alpha: np.ndarray, trimap: np.ndarray, settings: AlphaPostSettings
|
||||
) -> np.ndarray:
|
||||
cv2 = require_cv2()
|
||||
out = enforce_trimap(alpha, trimap)
|
||||
h, w = out.shape
|
||||
pixels = h * w
|
||||
min_component_area = max(4, int(round(pixels * settings.min_component_area_ratio)))
|
||||
fill_hole_area = max(4, int(round(pixels * settings.fill_hole_area_ratio)))
|
||||
|
||||
soft_fg = out > settings.alpha_floor
|
||||
cleaned_fg = _remove_small_components(soft_fg, min_component_area)
|
||||
out[~cleaned_fg & (trimap == 128)] = 0.0
|
||||
|
||||
solid = out >= settings.alpha_ceil
|
||||
inv = ~solid
|
||||
count, labels, stats, _ = cv2.connectedComponentsWithStats(inv.astype(np.uint8), 8)
|
||||
for label in range(1, count):
|
||||
area = stats[label, cv2.CC_STAT_AREA]
|
||||
touches_edge = (
|
||||
stats[label, cv2.CC_STAT_LEFT] == 0
|
||||
or stats[label, cv2.CC_STAT_TOP] == 0
|
||||
or stats[label, cv2.CC_STAT_LEFT] + stats[label, cv2.CC_STAT_WIDTH] >= w
|
||||
or stats[label, cv2.CC_STAT_TOP] + stats[label, cv2.CC_STAT_HEIGHT] >= h
|
||||
)
|
||||
if area <= fill_hole_area and not touches_edge:
|
||||
hole = labels == label
|
||||
out[hole & (trimap == 255)] = 1.0
|
||||
|
||||
return enforce_trimap(out, trimap)
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .deps import require_cv2
|
||||
from .settings import ChromaSettings
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BackgroundModel:
|
||||
rgb_center: tuple[float, float, float]
|
||||
rgb_sigma: tuple[float, float, float]
|
||||
hsv_center: tuple[float, float, float]
|
||||
hue_sigma: float
|
||||
lab_center: tuple[float, float, float]
|
||||
lab_sigma: float
|
||||
sample_count: int
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
def _smoothstep(x: np.ndarray, edge0: float, edge1: float) -> np.ndarray:
|
||||
t = np.clip((x - edge0) / max(edge1 - edge0, 1e-6), 0.0, 1.0)
|
||||
return t * t * (3.0 - 2.0 * t)
|
||||
|
||||
|
||||
def _hue_distance(hue: np.ndarray, center: float) -> np.ndarray:
|
||||
diff = np.abs(hue.astype(np.float32) - float(center))
|
||||
return np.minimum(diff, 180.0 - diff)
|
||||
|
||||
|
||||
def _border_mask(height: int, width: int, ratio: float) -> np.ndarray:
|
||||
border = max(8, int(round(max(height, width) * ratio)))
|
||||
border = min(border, height // 2, width // 2)
|
||||
mask = np.zeros((height, width), dtype=bool)
|
||||
mask[:border, :] = True
|
||||
mask[-border:, :] = True
|
||||
mask[:, :border] = True
|
||||
mask[:, -border:] = True
|
||||
return mask
|
||||
|
||||
|
||||
def convert_color_spaces(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
cv2 = require_cv2()
|
||||
rgb_f = rgb.astype(np.float32) / 255.0
|
||||
hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV).astype(np.float32)
|
||||
hsv[..., 1:] /= 255.0
|
||||
lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
return rgb_f, hsv, lab
|
||||
|
||||
|
||||
def initial_green_candidates(
|
||||
rgb_f: np.ndarray, hsv: np.ndarray, settings: ChromaSettings
|
||||
) -> np.ndarray:
|
||||
r = rgb_f[..., 0]
|
||||
g = rgb_f[..., 1]
|
||||
b = rgb_f[..., 2]
|
||||
dominance = g - np.maximum(r, b)
|
||||
hue = hsv[..., 0]
|
||||
sat = hsv[..., 1]
|
||||
return (
|
||||
(dominance >= settings.green_margin)
|
||||
& (g >= settings.min_green_value)
|
||||
& (sat >= settings.min_saturation)
|
||||
& (hue >= settings.hue_low)
|
||||
& (hue <= settings.hue_high)
|
||||
)
|
||||
|
||||
|
||||
def estimate_background_model(
|
||||
rgb: np.ndarray, settings: ChromaSettings
|
||||
) -> tuple[BackgroundModel, tuple[np.ndarray, np.ndarray, np.ndarray]]:
|
||||
rgb_f, hsv, lab = convert_color_spaces(rgb)
|
||||
candidates = initial_green_candidates(rgb_f, hsv, settings)
|
||||
h, w = candidates.shape
|
||||
border_candidates = candidates & _border_mask(h, w, settings.border_ratio)
|
||||
|
||||
sample_mask = border_candidates
|
||||
if int(sample_mask.sum()) < settings.min_samples:
|
||||
sample_mask = candidates
|
||||
if int(sample_mask.sum()) < max(64, settings.min_samples // 16):
|
||||
# Last-resort fallback: choose pixels with strongest green dominance.
|
||||
dominance = rgb_f[..., 1] - np.maximum(rgb_f[..., 0], rgb_f[..., 2])
|
||||
cutoff = np.percentile(dominance, 90.0)
|
||||
sample_mask = dominance >= cutoff
|
||||
|
||||
rgb_samples = rgb_f[sample_mask]
|
||||
hsv_samples = hsv[sample_mask]
|
||||
lab_samples = lab[sample_mask]
|
||||
|
||||
rgb_center = np.median(rgb_samples, axis=0)
|
||||
rgb_sigma = np.maximum(
|
||||
np.percentile(np.abs(rgb_samples - rgb_center), 75, axis=0) * 1.4826,
|
||||
settings.rgb_sigma_min,
|
||||
)
|
||||
hsv_center = np.median(hsv_samples, axis=0)
|
||||
hue_center = float(hsv_center[0])
|
||||
hue_dists = _hue_distance(hsv_samples[:, 0], hue_center)
|
||||
hue_sigma = max(
|
||||
float(np.percentile(hue_dists, 75) * 1.4826), settings.hue_sigma_min
|
||||
)
|
||||
lab_center = np.median(lab_samples, axis=0)
|
||||
lab_dists = np.linalg.norm(lab_samples - lab_center, axis=1)
|
||||
lab_sigma = max(
|
||||
float(np.percentile(lab_dists, 75) * 1.4826), settings.lab_sigma_min
|
||||
)
|
||||
|
||||
model = BackgroundModel(
|
||||
rgb_center=tuple(float(x) for x in rgb_center),
|
||||
rgb_sigma=tuple(float(x) for x in rgb_sigma),
|
||||
hsv_center=tuple(float(x) for x in hsv_center),
|
||||
hue_sigma=float(hue_sigma),
|
||||
lab_center=tuple(float(x) for x in lab_center),
|
||||
lab_sigma=float(lab_sigma),
|
||||
sample_count=int(sample_mask.sum()),
|
||||
)
|
||||
return model, (rgb_f, hsv, lab)
|
||||
|
||||
|
||||
def compute_bg_confidence(
|
||||
rgb: np.ndarray,
|
||||
model: BackgroundModel | None = None,
|
||||
settings: ChromaSettings | None = None,
|
||||
) -> tuple[np.ndarray, BackgroundModel]:
|
||||
settings = settings or ChromaSettings()
|
||||
if model is None:
|
||||
model, spaces = estimate_background_model(rgb, settings)
|
||||
else:
|
||||
spaces = convert_color_spaces(rgb)
|
||||
rgb_f, hsv, lab = spaces
|
||||
|
||||
r = rgb_f[..., 0]
|
||||
g = rgb_f[..., 1]
|
||||
b = rgb_f[..., 2]
|
||||
dominance = g - np.maximum(r, b)
|
||||
hue = hsv[..., 0]
|
||||
sat = hsv[..., 1]
|
||||
|
||||
hue_conf = np.exp(-0.5 * (_hue_distance(hue, model.hsv_center[0]) / model.hue_sigma) ** 2)
|
||||
sat_conf = _smoothstep(sat, max(0.05, model.hsv_center[1] * 0.45), max(0.2, model.hsv_center[1] * 0.85))
|
||||
dominance_conf = _smoothstep(dominance, settings.green_margin * 0.4, settings.green_margin * 1.6)
|
||||
|
||||
lab_center = np.asarray(model.lab_center, dtype=np.float32)
|
||||
lab_dist = np.linalg.norm(lab - lab_center, axis=2)
|
||||
lab_conf = np.exp(-0.5 * (lab_dist / model.lab_sigma) ** 2)
|
||||
|
||||
rgb_center = np.asarray(model.rgb_center, dtype=np.float32)
|
||||
rgb_sigma = np.asarray(model.rgb_sigma, dtype=np.float32)
|
||||
rgb_dist = np.linalg.norm((rgb_f - rgb_center) / rgb_sigma, axis=2)
|
||||
rgb_conf = np.exp(-0.5 * (rgb_dist / 2.5) ** 2)
|
||||
|
||||
conf = hue_conf * sat_conf * dominance_conf * np.maximum(lab_conf, rgb_conf * 0.85)
|
||||
return np.clip(conf, 0.0, 1.0).astype(np.float32), model
|
||||
+108
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
from rich import print
|
||||
|
||||
from .pipeline import MattingPipeline, run_image
|
||||
from .settings import (
|
||||
AlphaPostSettings,
|
||||
ChromaSettings,
|
||||
DespillSettings,
|
||||
ModelSettings,
|
||||
PipelineSettings,
|
||||
TrimapSettings,
|
||||
)
|
||||
|
||||
app = typer.Typer(help="Offline green screen character matting.")
|
||||
|
||||
|
||||
def _settings(
|
||||
model_name: str,
|
||||
device: str,
|
||||
matting_method: str,
|
||||
fallback_to_chroma_alpha: bool,
|
||||
sure_bg_threshold: float,
|
||||
sure_fg_threshold: float,
|
||||
unknown_radius_ratio: float,
|
||||
fg_safe_radius_ratio: float,
|
||||
despill: bool,
|
||||
) -> PipelineSettings:
|
||||
return PipelineSettings(
|
||||
chroma=ChromaSettings(),
|
||||
trimap=TrimapSettings(
|
||||
sure_bg_threshold=sure_bg_threshold,
|
||||
sure_fg_threshold=sure_fg_threshold,
|
||||
unknown_radius_ratio=unknown_radius_ratio,
|
||||
fg_safe_radius_ratio=fg_safe_radius_ratio,
|
||||
),
|
||||
alpha_post=AlphaPostSettings(),
|
||||
despill=DespillSettings(enabled=despill),
|
||||
model=ModelSettings(
|
||||
model_name=model_name,
|
||||
device=device,
|
||||
matting_method=matting_method,
|
||||
fallback_to_chroma_alpha=fallback_to_chroma_alpha,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def main(
|
||||
input: Path | None = typer.Option(None, "--input", "-i", exists=True, file_okay=True, dir_okay=False),
|
||||
output: Path | None = typer.Option(None, "--output", "-o", file_okay=True, dir_okay=False),
|
||||
input_dir: Path | None = typer.Option(None, "--input-dir", exists=True, file_okay=False, dir_okay=True),
|
||||
output_dir: Path | None = typer.Option(None, "--output-dir", file_okay=False, dir_okay=True),
|
||||
debug_dir: Path | None = typer.Option(None, "--debug-dir", file_okay=False, dir_okay=True),
|
||||
model_name: str = typer.Option("hustvl/vitmatte-small-composition-1k", "--model-name"),
|
||||
device: str = typer.Option("cuda", "--device"),
|
||||
matting_method: str = typer.Option("vitmatte", "--matting-method"),
|
||||
fallback_to_chroma_alpha: bool = typer.Option(False, "--fallback-to-chroma-alpha/--no-fallback-to-chroma-alpha"),
|
||||
sure_bg_threshold: float = typer.Option(0.92, "--sure-bg-threshold", min=0.0, max=1.0),
|
||||
sure_fg_threshold: float = typer.Option(0.12, "--sure-fg-threshold", min=0.0, max=1.0),
|
||||
unknown_radius_ratio: float = typer.Option(0.012, "--unknown-radius-ratio", min=0.0),
|
||||
fg_safe_radius_ratio: float = typer.Option(0.006, "--fg-safe-radius-ratio", min=0.0),
|
||||
despill: bool = typer.Option(True, "--despill/--no-despill"),
|
||||
) -> None:
|
||||
settings = _settings(
|
||||
model_name=model_name,
|
||||
device=device,
|
||||
matting_method=matting_method,
|
||||
fallback_to_chroma_alpha=fallback_to_chroma_alpha,
|
||||
sure_bg_threshold=sure_bg_threshold,
|
||||
sure_fg_threshold=sure_fg_threshold,
|
||||
unknown_radius_ratio=unknown_radius_ratio,
|
||||
fg_safe_radius_ratio=fg_safe_radius_ratio,
|
||||
despill=despill,
|
||||
)
|
||||
|
||||
try:
|
||||
if input_dir is not None:
|
||||
if output_dir is None:
|
||||
raise typer.BadParameter("--output-dir is required when --input-dir is used")
|
||||
images = sorted(
|
||||
p for p in input_dir.iterdir() if p.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}
|
||||
)
|
||||
if not images:
|
||||
raise typer.BadParameter(f"No supported images found in {input_dir}")
|
||||
pipeline = MattingPipeline(settings)
|
||||
for image_path in images:
|
||||
out_path = output_dir / f"{image_path.stem}_rgba.png"
|
||||
dbg = None if debug_dir is None else debug_dir / image_path.stem
|
||||
result = pipeline.run_image(image_path, out_path, dbg)
|
||||
print(f"[green]wrote[/green] {result['output']}")
|
||||
return
|
||||
|
||||
if input is None or output is None:
|
||||
raise typer.BadParameter("Use either --input/--output or --input-dir/--output-dir")
|
||||
|
||||
result = run_image(input, output, debug_dir, settings)
|
||||
print(f"[green]wrote[/green] {result['output']}")
|
||||
except RuntimeError as exc:
|
||||
print(f"[red]error:[/red] {exc}")
|
||||
raise typer.Exit(code=1) from exc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
@@ -0,0 +1,16 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def require_module(module_name: str, install_hint: str):
|
||||
try:
|
||||
return __import__(module_name)
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name == module_name:
|
||||
raise RuntimeError(
|
||||
f"Missing dependency '{module_name}'. Install it with: {install_hint}"
|
||||
) from exc
|
||||
raise
|
||||
|
||||
|
||||
def require_cv2():
|
||||
return require_module("cv2", "pip install opencv-python")
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .chroma import BackgroundModel
|
||||
from .settings import DespillSettings
|
||||
|
||||
|
||||
def despill_green(
|
||||
rgb: np.ndarray,
|
||||
alpha: np.ndarray,
|
||||
bg_confidence: np.ndarray,
|
||||
model: BackgroundModel,
|
||||
settings: DespillSettings,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
rgb_f = rgb.astype(np.float32) / 255.0
|
||||
if not settings.enabled:
|
||||
return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
|
||||
|
||||
edge = (alpha > settings.edge_low) & (alpha < settings.edge_high)
|
||||
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)
|
||||
|
||||
out = rgb_f.copy()
|
||||
out[..., 1] = g - excess * mask * settings.strength
|
||||
|
||||
# Very light compensation toward the non-green channels to avoid gray fringes.
|
||||
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)
|
||||
|
||||
return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def load_rgb(path: str | Path) -> np.ndarray:
|
||||
image = Image.open(path).convert("RGB")
|
||||
return np.asarray(image, dtype=np.uint8)
|
||||
|
||||
|
||||
def save_rgb(path: str | Path, rgb: np.ndarray) -> None:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
Image.fromarray(np.clip(rgb, 0, 255).astype(np.uint8), mode="RGB").save(path)
|
||||
|
||||
|
||||
def save_gray(path: str | Path, gray: np.ndarray) -> None:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if gray.dtype != np.uint8:
|
||||
gray = np.clip(gray * 255.0, 0, 255).astype(np.uint8)
|
||||
Image.fromarray(gray, mode="L").save(path)
|
||||
|
||||
|
||||
def save_rgba(path: str | Path, rgb: np.ndarray, alpha: np.ndarray) -> None:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
alpha_u8 = np.clip(alpha * 255.0, 0, 255).astype(np.uint8)
|
||||
rgba = np.dstack([np.clip(rgb, 0, 255).astype(np.uint8), alpha_u8])
|
||||
Image.fromarray(rgba, mode="RGBA").save(path)
|
||||
|
||||
|
||||
def write_text(path: str | Path, text: str) -> None:
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(text, encoding="utf-8")
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .alpha_post import clean_alpha, enforce_trimap
|
||||
from .chroma import compute_bg_confidence
|
||||
from .despill import despill_green
|
||||
from .io import load_rgb, save_gray, save_rgba, write_text
|
||||
from .qa import BACKGROUND_COLORS, composite, make_qa_grid, save_previews
|
||||
from .settings import PipelineSettings
|
||||
from .trimap import generate_trimap, trimap_to_alpha_seed
|
||||
from .vitmatte_infer import ViTMatteRunner
|
||||
|
||||
|
||||
class MattingPipeline:
|
||||
def __init__(self, settings: PipelineSettings):
|
||||
self.settings = settings
|
||||
self._runner: ViTMatteRunner | None = None
|
||||
|
||||
def _predict_alpha(self, rgb: np.ndarray, trimap: np.ndarray, bg_confidence: np.ndarray) -> tuple[np.ndarray, str]:
|
||||
if self.settings.model.matting_method == "chroma":
|
||||
return trimap_to_alpha_seed(trimap, bg_confidence), "chroma"
|
||||
if self.settings.model.matting_method != "vitmatte":
|
||||
raise RuntimeError(
|
||||
f"Unsupported matting method '{self.settings.model.matting_method}'. "
|
||||
"Use 'vitmatte' or 'chroma'."
|
||||
)
|
||||
|
||||
try:
|
||||
if self._runner is None:
|
||||
self._runner = ViTMatteRunner(self.settings.model)
|
||||
return self._runner.predict(rgb, trimap), "vitmatte"
|
||||
except RuntimeError:
|
||||
if not self.settings.model.fallback_to_chroma_alpha:
|
||||
raise
|
||||
return trimap_to_alpha_seed(trimap, bg_confidence), "chroma_fallback"
|
||||
|
||||
def run_image(
|
||||
self,
|
||||
input_path: str | Path,
|
||||
output_path: str | Path,
|
||||
debug_dir: str | Path | None,
|
||||
) -> dict:
|
||||
return _run_image(input_path, output_path, debug_dir, self)
|
||||
|
||||
|
||||
def run_image(
|
||||
input_path: str | Path,
|
||||
output_path: str | Path,
|
||||
debug_dir: str | Path | None,
|
||||
settings: PipelineSettings,
|
||||
) -> dict:
|
||||
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:
|
||||
settings = pipeline.settings
|
||||
rgb = load_rgb(input_path)
|
||||
bg_confidence, model = compute_bg_confidence(rgb, settings=settings.chroma)
|
||||
trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap)
|
||||
alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence)
|
||||
|
||||
alpha = enforce_trimap(alpha, trimap)
|
||||
alpha = clean_alpha(alpha, trimap, settings.alpha_post)
|
||||
corrected_rgb, despill_mask = despill_green(
|
||||
rgb, alpha, bg_confidence, model, settings.despill
|
||||
)
|
||||
save_rgba(output_path, corrected_rgb, alpha)
|
||||
|
||||
result = {
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"alpha_source": alpha_source,
|
||||
"background_model": model.to_dict(),
|
||||
"trimap": trimap_stats,
|
||||
"settings": asdict(settings),
|
||||
}
|
||||
|
||||
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)
|
||||
save_gray(debug / "alpha.png", alpha)
|
||||
save_gray(debug / "despill_mask.png", despill_mask)
|
||||
save_previews(debug, corrected_rgb, 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),
|
||||
}
|
||||
for name, color in BACKGROUND_COLORS.items():
|
||||
tiles[f"preview {name}"] = composite(corrected_rgb, alpha, color)
|
||||
make_qa_grid(debug, tiles)
|
||||
write_text(debug / "metadata.json", json.dumps(result, indent=2, ensure_ascii=False))
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
BACKGROUND_COLORS: dict[str, tuple[int, int, int]] = {
|
||||
"black": (0, 0, 0),
|
||||
"white": (255, 255, 255),
|
||||
"gray": (128, 128, 128),
|
||||
"red": (220, 32, 32),
|
||||
"blue": (32, 96, 220),
|
||||
}
|
||||
|
||||
|
||||
def composite(rgb: np.ndarray, alpha: np.ndarray, color: tuple[int, int, int]) -> np.ndarray:
|
||||
fg = rgb.astype(np.float32)
|
||||
bg = np.zeros_like(fg)
|
||||
bg[..., 0] = color[0]
|
||||
bg[..., 1] = color[1]
|
||||
bg[..., 2] = color[2]
|
||||
a = alpha[..., None].astype(np.float32)
|
||||
return np.clip(fg * a + bg * (1.0 - a), 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def save_previews(debug_dir: str | Path, rgb: np.ndarray, alpha: np.ndarray) -> dict[str, Path]:
|
||||
debug = Path(debug_dir)
|
||||
debug.mkdir(parents=True, exist_ok=True)
|
||||
paths: dict[str, Path] = {}
|
||||
for name, color in BACKGROUND_COLORS.items():
|
||||
path = debug / f"preview_{name}.png"
|
||||
Image.fromarray(composite(rgb, alpha, color), mode="RGB").save(path)
|
||||
paths[name] = path
|
||||
return paths
|
||||
|
||||
|
||||
def make_qa_grid(debug_dir: str | Path, tiles: dict[str, np.ndarray]) -> Path:
|
||||
debug = Path(debug_dir)
|
||||
debug.mkdir(parents=True, exist_ok=True)
|
||||
tile_items = list(tiles.items())
|
||||
if not tile_items:
|
||||
raise ValueError("No QA tiles provided")
|
||||
|
||||
thumb_w = 320
|
||||
label_h = 28
|
||||
padding = 10
|
||||
cols = min(3, len(tile_items))
|
||||
rows = int(np.ceil(len(tile_items) / cols))
|
||||
first_h, first_w = tile_items[0][1].shape[:2]
|
||||
thumb_h = max(180, int(round(thumb_w * first_h / max(first_w, 1))))
|
||||
|
||||
canvas = Image.new(
|
||||
"RGB",
|
||||
(
|
||||
cols * thumb_w + (cols + 1) * padding,
|
||||
rows * (thumb_h + label_h) + (rows + 1) * padding,
|
||||
),
|
||||
(32, 32, 32),
|
||||
)
|
||||
draw = ImageDraw.Draw(canvas)
|
||||
|
||||
for idx, (name, arr) in enumerate(tile_items):
|
||||
row = idx // cols
|
||||
col = idx % cols
|
||||
x = padding + col * (thumb_w + padding)
|
||||
y = padding + row * (thumb_h + label_h + padding)
|
||||
img = Image.fromarray(arr.astype(np.uint8), mode="RGB")
|
||||
img.thumbnail((thumb_w, thumb_h), Image.Resampling.LANCZOS)
|
||||
frame = Image.new("RGB", (thumb_w, thumb_h), (240, 240, 240))
|
||||
frame.paste(img, ((thumb_w - img.width) // 2, (thumb_h - img.height) // 2))
|
||||
canvas.paste(frame, (x, y + label_h))
|
||||
draw.text((x, y), name, fill=(245, 245, 245))
|
||||
|
||||
path = debug / "qa_grid.png"
|
||||
canvas.save(path)
|
||||
return path
|
||||
@@ -0,0 +1,61 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ChromaSettings:
|
||||
green_margin: float = 0.08
|
||||
min_green_value: float = 0.22
|
||||
min_saturation: float = 0.20
|
||||
hue_low: float = 35.0
|
||||
hue_high: float = 95.0
|
||||
border_ratio: float = 0.04
|
||||
min_samples: int = 2048
|
||||
hue_sigma_min: float = 7.0
|
||||
lab_sigma_min: float = 10.0
|
||||
rgb_sigma_min: float = 0.08
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TrimapSettings:
|
||||
sure_bg_threshold: float = 0.92
|
||||
sure_fg_threshold: float = 0.12
|
||||
unknown_radius_ratio: float = 0.012
|
||||
fg_safe_radius_ratio: float = 0.006
|
||||
min_unknown_radius: int = 4
|
||||
min_fg_safe_radius: int = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AlphaPostSettings:
|
||||
min_component_area_ratio: float = 0.00001
|
||||
fill_hole_area_ratio: float = 0.00002
|
||||
alpha_floor: float = 0.002
|
||||
alpha_ceil: float = 0.998
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelSettings:
|
||||
model_name: str = "hustvl/vitmatte-small-composition-1k"
|
||||
device: str = "cuda"
|
||||
matting_method: str = "vitmatte"
|
||||
fallback_to_chroma_alpha: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineSettings:
|
||||
chroma: ChromaSettings = ChromaSettings()
|
||||
trimap: TrimapSettings = TrimapSettings()
|
||||
alpha_post: AlphaPostSettings = AlphaPostSettings()
|
||||
despill: DespillSettings = DespillSettings()
|
||||
model: ModelSettings = ModelSettings()
|
||||
@@ -0,0 +1,60 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .deps import require_cv2
|
||||
from .settings import TrimapSettings
|
||||
|
||||
|
||||
def radius_from_ratio(shape: tuple[int, int], ratio: float, minimum: int) -> int:
|
||||
return max(minimum, int(round(max(shape) * ratio)))
|
||||
|
||||
|
||||
def elliptical_kernel(radius: int) -> np.ndarray:
|
||||
cv2 = require_cv2()
|
||||
size = radius * 2 + 1
|
||||
return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size))
|
||||
|
||||
|
||||
def generate_trimap(
|
||||
bg_confidence: np.ndarray, settings: TrimapSettings
|
||||
) -> tuple[np.ndarray, dict[str, int]]:
|
||||
cv2 = require_cv2()
|
||||
shape = bg_confidence.shape
|
||||
unknown_radius = radius_from_ratio(
|
||||
shape, settings.unknown_radius_ratio, settings.min_unknown_radius
|
||||
)
|
||||
fg_safe_radius = radius_from_ratio(
|
||||
shape, settings.fg_safe_radius_ratio, settings.min_fg_safe_radius
|
||||
)
|
||||
|
||||
sure_bg = bg_confidence >= settings.sure_bg_threshold
|
||||
low_bg = bg_confidence <= settings.sure_fg_threshold
|
||||
|
||||
bg_u8 = sure_bg.astype(np.uint8)
|
||||
unknown_band = cv2.dilate(bg_u8, elliptical_kernel(unknown_radius)).astype(bool)
|
||||
fg_safe = ~cv2.dilate(bg_u8, elliptical_kernel(fg_safe_radius)).astype(bool)
|
||||
sure_fg = low_bg & fg_safe
|
||||
|
||||
trimap = np.full(shape, 128, dtype=np.uint8)
|
||||
trimap[sure_bg] = 0
|
||||
trimap[sure_fg] = 255
|
||||
|
||||
# Keep a protective unknown band around all sure background, including holes.
|
||||
trimap[unknown_band & ~sure_bg & ~sure_fg] = 128
|
||||
|
||||
stats = {
|
||||
"sure_bg_pixels": int((trimap == 0).sum()),
|
||||
"unknown_pixels": int((trimap == 128).sum()),
|
||||
"sure_fg_pixels": int((trimap == 255).sum()),
|
||||
"unknown_radius": int(unknown_radius),
|
||||
"fg_safe_radius": int(fg_safe_radius),
|
||||
}
|
||||
return trimap, stats
|
||||
|
||||
|
||||
def trimap_to_alpha_seed(trimap: np.ndarray, bg_confidence: np.ndarray) -> np.ndarray:
|
||||
alpha = np.clip(1.0 - bg_confidence, 0.0, 1.0).astype(np.float32)
|
||||
alpha[trimap == 0] = 0.0
|
||||
alpha[trimap == 255] = 1.0
|
||||
return alpha
|
||||
@@ -0,0 +1,47 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from .settings import ModelSettings
|
||||
|
||||
|
||||
class ViTMatteRunner:
|
||||
def __init__(self, settings: ModelSettings):
|
||||
try:
|
||||
import torch
|
||||
from transformers import VitMatteForImageMatting, VitMatteImageProcessor
|
||||
except ModuleNotFoundError as exc:
|
||||
raise RuntimeError(
|
||||
"Missing ViTMatte dependencies. Install them with: "
|
||||
"pip install torch transformers accelerate safetensors"
|
||||
) from exc
|
||||
|
||||
self.torch = torch
|
||||
self.processor = VitMatteImageProcessor.from_pretrained(settings.model_name)
|
||||
self.model = VitMatteForImageMatting.from_pretrained(settings.model_name)
|
||||
self.device = self._resolve_device(settings.device)
|
||||
self.model.to(self.device)
|
||||
self.model.eval()
|
||||
|
||||
def _resolve_device(self, requested: str):
|
||||
if requested == "cuda" and not self.torch.cuda.is_available():
|
||||
raise RuntimeError(
|
||||
"CUDA was requested but is not available. Use --device cpu or install a CUDA-enabled PyTorch build."
|
||||
)
|
||||
return self.torch.device(requested)
|
||||
|
||||
def predict(self, rgb: np.ndarray, trimap: np.ndarray) -> np.ndarray:
|
||||
image = Image.fromarray(rgb.astype(np.uint8), mode="RGB")
|
||||
trimap_image = Image.fromarray(trimap.astype(np.uint8), mode="L")
|
||||
inputs = self.processor(images=image, trimaps=trimap_image, return_tensors="pt")
|
||||
inputs = {key: value.to(self.device) for key, value in inputs.items()}
|
||||
with self.torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
alpha = outputs.alphas[0, 0].detach().float().cpu().numpy()
|
||||
if alpha.shape != trimap.shape:
|
||||
alpha_img = Image.fromarray(np.clip(alpha * 255.0, 0, 255).astype(np.uint8), mode="L")
|
||||
alpha_img = alpha_img.resize((trimap.shape[1], trimap.shape[0]), Image.Resampling.BILINEAR)
|
||||
alpha = np.asarray(alpha_img, dtype=np.float32) / 255.0
|
||||
return np.clip(alpha, 0.0, 1.0).astype(np.float32)
|
||||
@@ -0,0 +1,17 @@
|
||||
model:
|
||||
model_name: hustvl/vitmatte-small-composition-1k
|
||||
device: cuda
|
||||
matting_method: vitmatte
|
||||
fallback_to_chroma_alpha: false
|
||||
|
||||
trimap:
|
||||
sure_bg_threshold: 0.92
|
||||
sure_fg_threshold: 0.12
|
||||
unknown_radius_ratio: 0.012
|
||||
fg_safe_radius_ratio: 0.006
|
||||
|
||||
despill:
|
||||
enabled: true
|
||||
edge_low: 0.02
|
||||
edge_high: 0.98
|
||||
strength: 0.75
|
||||
@@ -129,6 +129,7 @@ python -m bgfilter.cli `
|
||||
```text
|
||||
--model-name hustvl/vitmatte-small-composition-1k
|
||||
--device cuda
|
||||
--matting-method vitmatte
|
||||
--sure-bg-threshold 0.92
|
||||
--sure-fg-threshold 0.12
|
||||
--unknown-radius-ratio 0.012
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
opencv-python
|
||||
pillow
|
||||
numpy
|
||||
scipy
|
||||
scikit-image
|
||||
torch
|
||||
transformers
|
||||
accelerate
|
||||
safetensors
|
||||
huggingface_hub
|
||||
tqdm
|
||||
typer
|
||||
rich
|
||||
Reference in New Issue
Block a user