Spaces:
Sleeping
Sleeping
File size: 2,246 Bytes
419fc11 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Any, Dict
import torch
from infer.pipeline import load_checkpoint, predict_image
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run single-image inference for baseline or GNN checkpoints.")
parser.add_argument("--model-type", type=str, choices=["auto", "baseline", "gnn"], default="auto")
parser.add_argument(
"--checkpoint",
type=Path,
default=Path("outputs/smoke_baseline/best.pt"),
help="Path to checkpoint (.pt).",
)
parser.add_argument(
"--image",
type=Path,
required=True,
help="Path to input image.",
)
parser.add_argument("--top-k", type=int, default=5)
parser.add_argument("--threshold", type=float, default=0.5)
parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"])
parser.add_argument("--benchmark", action="store_true", help="Print simple latency timing.")
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.device == "auto":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
device = torch.device(args.device)
loaded = load_checkpoint(args.checkpoint, device=device)
if args.benchmark:
import time
start = time.time()
result = predict_image(
image_path=args.image,
loaded=loaded,
top_k=args.top_k,
threshold=args.threshold,
)
elapsed = time.time() - start
else:
result = predict_image(
image_path=args.image,
loaded=loaded,
top_k=args.top_k,
threshold=args.threshold,
)
elapsed = None
payload: Dict[str, Any] = {
"image": str(args.image),
"top_k": args.top_k,
"threshold": args.threshold,
"labels": result.labels,
"scores": result.scores,
"positives": result.positives,
}
if elapsed is not None:
payload["latency_sec"] = elapsed
print(json.dumps(payload, indent=2))
if __name__ == "__main__":
main()
|