add open-vocab YOLO-World engine with CLIP text encoder

- Add OpenVocab engine: YOLO-World ONNX detection with runtime
  vocabulary via CLIP text embeddings (two-input model support)
- Add header-only C++ CLIP BPE tokenizer (no external deps)
- Add detect(cv::Mat) overload to OpenVocab for camera-free testing
- Add standalone test binary (test_open_vocab) using test/assets images
- Add export scripts for YOLO-World and CLIP text encoder (uv/PEP 723)
- Add installation scripts for Linux x64, ARM (Jetson), and Windows
- Update CMakeLists to use find_package(onnxruntime) instead of FetchContent
- Update README and AGENTS with env vars, setup workflow, and test usage
This commit is contained in:
cristhian aguilera
2026-03-28 10:41:11 -03:00
parent 6e6dc8ef27
commit fd6993e294
18 changed files with 1611 additions and 37 deletions
+124
View File
@@ -0,0 +1,124 @@
# /// script
# requires-python = ">=3.10"
# dependencies = ["ultralytics>=8.0"]
# ///
"""Export a YOLOv8-worldv2 model to ONNX with N open-vocabulary class slots.
The exported model has two inputs:
images (1, 3, imgsz, imgsz) float32
text_feats (1, N, 512) float32 ← CLIP embeddings, set at runtime
N is fixed per export. Use export_clip_text_encoder.py to generate the text
encoder and vocab files used to compute embeddings at runtime in the C++ node.
Usage:
uv run scripts/export_yoloworld.py
uv run scripts/export_yoloworld.py --model m --num-classes 20
uv run scripts/export_yoloworld.py --model l --num-classes 5 --output-dir models/
"""
import argparse
import shutil
import sys
from pathlib import Path
MODELS = {
"s": "yolov8s-worldv2",
"m": "yolov8m-worldv2",
"l": "yolov8l-worldv2",
"x": "yolov8x-worldv2",
}
def parse_args():
p = argparse.ArgumentParser(
description="Export YOLOv8-worldv2 to ONNX with open-vocabulary text inputs"
)
p.add_argument(
"--model", "-m",
choices=list(MODELS),
default="s",
help="Model size: s (~14 MB), m (~30 MB), l (~50 MB), x (~100 MB). Default: s",
)
p.add_argument(
"--num-classes", "-n",
type=int,
default=10,
help="Number of concurrent class slots (must match number of labels at runtime). Default: 10",
)
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_inputs(onnx_path: Path, expected_inputs: int):
try:
import onnx
model = onnx.load(str(onnx_path))
actual = len(model.graph.input)
if actual != expected_inputs:
print(f"\nWarning: expected {expected_inputs} inputs, got {actual}.")
print("The visual model may have baked text embeddings (ultralytics version dependent).")
print("In that case, use a fixed vocabulary instead of the text encoder.")
else:
inputs = [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim])
for i in model.graph.input]
print(f" Inputs: {inputs}")
except ImportError:
print(" (install 'onnx' to verify model inputs)")
def main():
args = parse_args()
try:
from ultralytics import YOLO
except ImportError:
print("Error: ultralytics not installed. Run: uv run scripts/export_yoloworld.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}_n{args.num_classes}_i{args.imgsz}.onnx"
dst = out_dir / out_name
print(f"Model: {model_name}.pt")
print(f"Class slots: {args.num_classes}")
print(f"Input size: {args.imgsz}")
print(f"Output: {dst}")
print()
model = YOLO(f"{model_name}.pt")
# Set N dummy class slots — actual label meanings are provided at runtime via text_feats
model.set_classes(["object"] * args.num_classes)
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_inputs(dst, expected_inputs=2)
print(f"\nExported: {dst.resolve()}")
print(f"\nNext step — export the CLIP text encoder (once, reused for all model sizes):")
print(f" uv run scripts/export_clip_text_encoder.py --output-dir models/")
print(f"\nRun with:")
print(f" ZED_ENGINE=open_vocab \\")
print(f" ZED_ONNX_MODEL_PATH={dst.resolve()} \\")
print(f" ZED_ONNX_TEXT_ENCODER_PATH=models/clip_text_encoder.onnx \\")
print(f' ZED_ONNX_LABELS="person,car,bag,... (exactly {args.num_classes} labels)" \\')
print(f" dora start dataflow.yml")
if __name__ == "__main__":
main()