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:
@@ -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()
|
||||
Reference in New Issue
Block a user