refactor open_vocab into yolo_onnx and yoloworld_onnx engines

Split the monolithic open_vocab engine into two focused engines: yolo_onnx (closed-vocabulary, any single-input YOLO ONNX) and yoloworld_onnx (open-vocabulary YOLO-World with CLIP encoder). Shared NMS/letterbox/fuse helpers extracted to yolo_helpers.h. Add export_yolo.py script and separate test binaries for each engine.
2026-03-30 04:44:25 -03:00
parent 748b0c3665
commit 44eec29b83
14 changed files with 1467 additions and 825 deletions
@@ -0,0 +1,111 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["ultralytics>=8.3"]
+# ///
+"""Export a standard YOLOv11 model to ONNX for closed-vocabulary detection.
+
+The exported model has one input:
+  images  (1, 3, imgsz, imgsz)  float32
+
+Class names are fixed at training time. Provide them at runtime via
+ZED_ONNX_LABELS or ZED_ONNX_LABELS_FILE.
+
+Usage:
+    uv run scripts/export_yolo.py
+    uv run scripts/export_yolo.py --model s --output-dir models/
+    uv run scripts/export_yolo.py --model m --imgsz 1280
+"""
+
+import argparse
+import shutil
+import sys
+from pathlib import Path
+
+MODELS = {
+    "n": "yolo11n",
+    "s": "yolo11s",
+    "m": "yolo11m",
+    "l": "yolo11l",
+    "x": "yolo11x",
+}
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Export YOLOv11 to ONNX for closed-vocabulary detection"
+    )
+    p.add_argument(
+        "--model", "-m",
+        choices=list(MODELS),
+        default="n",
+        help="Model size: n (~6 MB), s (~18 MB), m (~38 MB), l (~49 MB), x (~109 MB). Default: n",
+    )
+    p.add_argument(
+        "--imgsz",
+        type=int,
+        default=640,
+        help="Input image size in pixels (square). Default: 640",
+    )
+    p.add_argument(
+        "--output-dir", "-o",
+        default=".",
+        help="Output directory for the .onnx file. Default: current directory",
+    )
+    return p.parse_args()
+
+
+def verify_output(onnx_path: Path):
+    try:
+        import onnx
+        model = onnx.load(str(onnx_path))
+        inputs  = [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim])
+                   for i in model.graph.input]
+        outputs = [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim])
+                   for o in model.graph.output]
+        print(f"  Inputs:  {inputs}")
+        print(f"  Outputs: {outputs}")
+    except ImportError:
+        print("  (install 'onnx' to verify model structure)")
+
+
+def main():
+    args = parse_args()
+
+    try:
+        from ultralytics import YOLO
+    except ImportError:
+        print("Error: ultralytics not installed. Run: uv run scripts/export_yolo.py", file=sys.stderr)
+        sys.exit(1)
+
+    model_name = MODELS[args.model]
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_name = f"{model_name}_i{args.imgsz}.onnx"
+    dst = out_dir / out_name
+
+    print(f"Model:      {model_name}.pt")
+    print(f"Input size: {args.imgsz}")
+    print(f"Output:     {dst}")
+    print()
+
+    model = YOLO(f"{model_name}.pt")
+    exported = Path(model.export(format="onnx", imgsz=args.imgsz))
+
+    if exported.resolve() != dst.resolve():
+        shutil.move(str(exported), str(dst))
+
+    print(f"\nVerifying exported model...")
+    verify_output(dst)
+
+    print(f"\nExported: {dst.resolve()}")
+    print(f"\nRun with (COCO labels example):")
+    print(f"  ZED_ENGINE=onnx \\")
+    print(f"  ZED_ONNX_MODEL_PATH={dst.resolve()} \\")
+    print(f"  ZED_ONNX_LABELS_FILE=models/coco80.txt \\")
+    print(f"  dora start dataflow.yml")
+    print(f"\nOr with inline labels:")
+    print(f"  ZED_ONNX_LABELS=\"person,car,chair\" \\")
+
+
+if __name__ == "__main__":
+    main()