File size: 1,494 Bytes
0d0ec29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | #!/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()
|