scripts/example_inference.py

#!/usr/bin/env python3
"""Minimal paired-image inference example.

Run from the GitHub repository root after installing repository requirements.
"""

from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image

from src.engine.checkpoint import load_checkpoint
from src.engine.inference import predict_pair
from src.models.build import build_model
from src.utils.config import load_config
from src.utils.device import get_device


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="configs/active.yaml")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--image-a", required=True)
    parser.add_argument("--image-b", required=True)
    parser.add_argument("--output", default="prediction.png")
    parser.add_argument("--threshold", type=float, default=None)
    args = parser.parse_args()

    cfg = load_config(args.config)
    device = get_device(cfg)
    model = build_model(cfg).to(device)
    checkpoint = load_checkpoint(args.checkpoint, model)
    threshold = args.threshold
    if threshold is None:
        threshold = float(checkpoint.get("best_threshold", cfg.eval.threshold))
    model.eval()
    pred = predict_pair(model, Path(args.image_a), Path(args.image_b), cfg, device, threshold)
    Image.fromarray((pred.numpy() * 255).astype("uint8")).save(args.output)
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()