BGfilter/bgfilter/cli.py

from __future__ import annotations

from pathlib import Path

import typer
from rich import print

from .config import load_settings, override_settings
from .memtune import release_freed_memory
from .pipeline import MattingPipeline, run_image

app = typer.Typer(help="Offline flat-background character matting.")


@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),
    config: Path | None = typer.Option(None, "--config", exists=True, file_okay=True, dir_okay=False),
    model_name: str | None = typer.Option(None, "--model-name"),
    device: str | None = typer.Option(None, "--device"),
    precision: str | None = typer.Option(None, "--precision", help="Compute precision for all models (ViTMatte cast + BiRefNet autocast): fp32 (default) | bf16 (faster + halves matting activation memory; needs bf16-capable hardware, else falls back to fp32)"),
    screen_color: str | None = typer.Option(None, "--screen-color", help="Background colour prior as #RRGGBB (default: auto-detect the flat background colour)"),
    matting_method: str | None = typer.Option(None, "--matting-method"),
    fallback_to_chroma_alpha: bool | None = typer.Option(None, "--fallback-to-chroma-alpha/--no-fallback-to-chroma-alpha"),
    sure_bg_threshold: float | None = typer.Option(None, "--sure-bg-threshold", min=0.0, max=1.0),
    sure_fg_threshold: float | None = typer.Option(None, "--sure-fg-threshold", min=0.0, max=1.0),
    unknown_radius_ratio: float | None = typer.Option(None, "--unknown-radius-ratio", min=0.0),
    fg_safe_radius_ratio: float | None = typer.Option(None, "--fg-safe-radius-ratio", min=0.0),
    despill: bool | None = typer.Option(None, "--despill/--no-despill"),
    cross_check: bool | None = typer.Option(None, "--cross-check/--no-cross-check", help="Second-opinion veto of background-hued residue between hair strands (default: on; costs one extra model inference)"),
    cross_check_as_seg: bool | None = typer.Option(None, "--cross-check-as-seg/--no-cross-check-as-seg", help="Reuse the cross-check forward as the segmentation mask, skipping the primary seg model (default: off; saves ~20s on CPU; needs cross-check on)"),
    trimap_mode: str | None = typer.Option(None, "--trimap-mode", help="Trimap mode (segmentation pipeline): directional | seg | directional-hard-bg"),
    seg_backend: str | None = typer.Option(None, "--seg-backend", help="Segmentation backend: birefnet (default) | anime-seg"),
) -> None:
    try:
        settings = override_settings(
            load_settings(config),
            model_name=model_name,
            device=device,
            precision=precision,
            screen_color=screen_color,
            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,
            cross_check=cross_check,
            cross_check_as_seg=cross_check_as_seg,
            trimap_mode=trimap_mode,
            seg_backend=seg_backend,
        )

        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']}")
                # Free this image's activations back to the OS before the next one
                # so batch RSS tracks a single image, not the whole run.
                release_freed_memory()
            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
    except ValueError as exc:
        print(f"[red]config error:[/red] {exc}")
        raise typer.Exit(code=1) from exc


if __name__ == "__main__":
    app()