predict.py

#!/usr/bin/env python3
"""HARRIER YOLO26 single-shot prediction CLI.

Example:

    python scripts/predict.py --weights runs/best.pt \\
        --modality eo --source /path/to/image.jpg
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT / "src"))

from anima_harrier.data.indraeye import IndraEyeModality  # noqa: E402
from anima_harrier.inference import HarrierPredictor, PredictRequest  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="anima-harrier-predict")
    parser.add_argument("--weights", type=Path, required=True)
    parser.add_argument("--modality", choices=[m.value for m in IndraEyeModality], required=True)
    parser.add_argument("--source", required=True, help="Image/video path or stream URI.")
    parser.add_argument("--conf", type=float, default=0.25)
    parser.add_argument("--iou", type=float, default=0.45)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--backend", default="auto", choices=["auto", "cuda", "mlx", "cpu"])
    args = parser.parse_args(argv)

    predictor = HarrierPredictor(weights=args.weights, backend=args.backend)
    request = PredictRequest(
        modality=IndraEyeModality(args.modality),
        source=args.source,
        conf=args.conf,
        iou=args.iou,
        imgsz=args.imgsz,
    )
    response = predictor.predict(request)
    print(json.dumps(response.to_dict(), indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())