diff --git a/.gitignore b/.gitignore index 0ba96f3..a6f8414 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ Samples/ Outputs/ +__pycache__/ +*.py[cod] diff --git a/README.md b/README.md new file mode 100644 index 0000000..b9c7c90 --- /dev/null +++ b/README.md @@ -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. diff --git a/bgfilter/__init__.py b/bgfilter/__init__.py new file mode 100644 index 0000000..2e425e6 --- /dev/null +++ b/bgfilter/__init__.py @@ -0,0 +1,3 @@ +"""Green screen character matting pipeline.""" + +__version__ = "0.1.0" diff --git a/bgfilter/alpha_post.py b/bgfilter/alpha_post.py new file mode 100644 index 0000000..9e1f0b6 --- /dev/null +++ b/bgfilter/alpha_post.py @@ -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) diff --git a/bgfilter/chroma.py b/bgfilter/chroma.py new file mode 100644 index 0000000..3ddc77b --- /dev/null +++ b/bgfilter/chroma.py @@ -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 diff --git a/bgfilter/cli.py b/bgfilter/cli.py new file mode 100644 index 0000000..dfbb1a6 --- /dev/null +++ b/bgfilter/cli.py @@ -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() diff --git a/bgfilter/deps.py b/bgfilter/deps.py new file mode 100644 index 0000000..d79783b --- /dev/null +++ b/bgfilter/deps.py @@ -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") diff --git a/bgfilter/despill.py b/bgfilter/despill.py new file mode 100644 index 0000000..c38f875 --- /dev/null +++ b/bgfilter/despill.py @@ -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) diff --git a/bgfilter/io.py b/bgfilter/io.py new file mode 100644 index 0000000..b72f3de --- /dev/null +++ b/bgfilter/io.py @@ -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") diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py new file mode 100644 index 0000000..89ea2c7 --- /dev/null +++ b/bgfilter/pipeline.py @@ -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 diff --git a/bgfilter/qa.py b/bgfilter/qa.py new file mode 100644 index 0000000..df660a2 --- /dev/null +++ b/bgfilter/qa.py @@ -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 diff --git a/bgfilter/settings.py b/bgfilter/settings.py new file mode 100644 index 0000000..a31435b --- /dev/null +++ b/bgfilter/settings.py @@ -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() diff --git a/bgfilter/trimap.py b/bgfilter/trimap.py new file mode 100644 index 0000000..082e9f1 --- /dev/null +++ b/bgfilter/trimap.py @@ -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 diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py new file mode 100644 index 0000000..d33a265 --- /dev/null +++ b/bgfilter/vitmatte_infer.py @@ -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) diff --git a/configs/default.yaml b/configs/default.yaml new file mode 100644 index 0000000..a953a45 --- /dev/null +++ b/configs/default.yaml @@ -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 diff --git a/docs/green_screen_matting_workflow.md b/docs/green_screen_matting_workflow.md index c7e0b70..a92fcf3 100644 --- a/docs/green_screen_matting_workflow.md +++ b/docs/green_screen_matting_workflow.md @@ -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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a410475 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,13 @@ +opencv-python +pillow +numpy +scipy +scikit-image +torch +transformers +accelerate +safetensors +huggingface_hub +tqdm +typer +rich