File size: 1,406 Bytes
7fb0075 | 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 47 48 49 50 51 52 53 54 | #!/usr/bin/env python3
"""CLI: image path to mili score in [0, 1]. --backbone: local directory or Hugging Face model id."""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
from PIL import Image
from mili_score_inference.predict import load_model, predict_pil
def main() -> int:
p = argparse.ArgumentParser()
p.add_argument("--weights", type=Path, required=True)
p.add_argument(
"--backbone",
type=str,
required=True,
help="Local checkpoint directory or Hugging Face model id",
)
p.add_argument("--image", type=Path, required=True)
p.add_argument("--device", default="auto")
args = p.parse_args()
if args.device == "auto":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
device = torch.device(args.device)
bb = Path(args.backbone)
if bb.is_dir():
backbone_arg: Path | str = bb
local_bb = True
else:
backbone_arg = args.backbone
local_bb = False
model = load_model(
args.weights,
backbone_arg,
device=device,
local_backbone=local_bb,
)
score = predict_pil(model, Image.open(args.image), device=device)
print(f"score={score:.6f}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|