65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Single-image inference with CorridorKey MLX.
|
|
|
|
Usage:
|
|
uv run python scripts/infer.py --image input.png --hint alpha_hint.png
|
|
uv run python scripts/infer.py --image input.png --hint alpha_hint.png --output-dir results/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from corridorkey_mlx.inference.pipeline import (
|
|
DEFAULT_CHECKPOINT,
|
|
DEFAULT_IMG_SIZE,
|
|
infer_and_save,
|
|
load_model,
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="CorridorKey MLX inference")
|
|
parser.add_argument("--image", type=Path, required=True, help="RGB input image")
|
|
parser.add_argument("--hint", type=Path, required=True, help="Alpha hint (grayscale)")
|
|
parser.add_argument(
|
|
"--checkpoint",
|
|
type=Path,
|
|
default=DEFAULT_CHECKPOINT,
|
|
help="MLX safetensors checkpoint",
|
|
)
|
|
parser.add_argument("--img-size", type=int, default=DEFAULT_IMG_SIZE, help="Model input size")
|
|
parser.add_argument("--output-dir", type=Path, default=Path("output"), help="Output directory")
|
|
args = parser.parse_args()
|
|
|
|
if not args.image.exists():
|
|
print(f"Image not found: {args.image}")
|
|
raise SystemExit(1)
|
|
if not args.hint.exists():
|
|
print(f"Alpha hint not found: {args.hint}")
|
|
raise SystemExit(1)
|
|
if not args.checkpoint.exists():
|
|
print(f"Checkpoint not found: {args.checkpoint}")
|
|
print("Run: uv run python scripts/convert_weights.py")
|
|
raise SystemExit(1)
|
|
|
|
print(f"Loading model (img_size={args.img_size})...")
|
|
t0 = time.perf_counter()
|
|
model = load_model(args.checkpoint, args.img_size)
|
|
print(f" Model loaded in {time.perf_counter() - t0:.2f}s")
|
|
|
|
print(f"Running inference on {args.image}...")
|
|
t0 = time.perf_counter()
|
|
results = infer_and_save(model, args.image, args.hint, args.output_dir)
|
|
print(f" Inference + save in {time.perf_counter() - t0:.2f}s")
|
|
|
|
print(f"\nSaved: {args.output_dir / 'alpha.png'}, {args.output_dir / 'foreground.png'}")
|
|
print(f" Alpha shape: {results['alpha'].shape}")
|
|
print(f" Foreground shape: {results['foreground'].shape}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|