scorevision: push artifact

training/export_onnx.py

@@ -0,0 +1,65 @@
+"""
+Export a trained YOLOv11 .pt to ONNX with NMS baked in -> output shape [1, max_det, 6].
+
+The default chute miner.py decode path expects [1, N, 6] (NMS-baked). The king's own
+weights.onnx is raw-YOLO [1, 10, 8400] without NMS — our exported model does NOT match
+that, but our miner.py handles both formats. NMS-baked is preferred because:
+  - chute startup is faster (less graph work),
+  - the decode path is simpler (no fallback to raw YOLO numpy NMS),
+  - p95 latency is more stable (no Python-NMS spike on busy frames).
+
+Usage:
+  python export_onnx.py --weights runs/detect/crime_b/weights/best.pt --imgsz 1280 \
+                        --out ../weights.onnx
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+from pathlib import Path
+
+from ultralytics import YOLO
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--weights", required=True)
+    ap.add_argument("--imgsz", type=int, default=1280)
+    ap.add_argument("--out", default=None,
+                    help="destination path for weights.onnx (defaults next to source .pt)")
+    ap.add_argument("--max-det", type=int, default=300)
+    ap.add_argument("--opset", type=int, default=17)
+    ap.add_argument("--half", action="store_true",
+                    help="export FP16 (smaller, faster on CUDA, marginal accuracy delta)")
+    args = ap.parse_args()
+
+    weights = Path(args.weights).resolve()
+    model = YOLO(str(weights))
+
+    onnx_path = model.export(
+        format="onnx",
+        imgsz=args.imgsz,
+        nms=True,
+        max_det=args.max_det,
+        opset=args.opset,
+        dynamic=False,
+        half=args.half,
+        simplify=True,
+        agnostic_nms=False,   # class-aware NMS (must match miner.py expectations)
+        device="cpu",
+    )
+
+    src = Path(onnx_path)
+    if args.out:
+        dst = Path(args.out).resolve()
+        dst.parent.mkdir(parents=True, exist_ok=True)
+        shutil.move(str(src), str(dst))
+        print(f"[export] {src} -> {dst}")
+    else:
+        print(f"[export] {src}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())