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
@@ -0,0 +1,105 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["transformers>=4.30", "torch", "onnx"]
|
||||
# ///
|
||||
"""Export the CLIP ViT-B/32 text encoder to ONNX and save vocab files.
|
||||
|
||||
Compatible with all YOLOv8-worldv2 sizes (s/m/l/x) — run once and reuse.
|
||||
|
||||
The C++ node uses these three files to compute text embeddings at runtime:
|
||||
clip_text_encoder.onnx input: input_ids (1, 77) int64
|
||||
output: text_embeds (1, 512) float32 (L2-normalize in C++)
|
||||
clip_vocab.txt token → id mapping (tab-separated, one per line)
|
||||
clip_merges.txt BPE merge rules (space-separated pairs, one per line)
|
||||
|
||||
Usage:
|
||||
uv run scripts/export_clip_text_encoder.py
|
||||
uv run scripts/export_clip_text_encoder.py --output-dir models/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
CLIP_MODEL = "openai/clip-vit-base-patch32"
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Export CLIP ViT-B/32 text encoder to ONNX")
|
||||
p.add_argument(
|
||||
"--output-dir", "-o",
|
||||
default="models",
|
||||
help="Directory for output files. Default: models/",
|
||||
)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
try:
|
||||
import torch
|
||||
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
|
||||
except ImportError as e:
|
||||
print(f"Error: {e}\nRun: uv run scripts/export_clip_text_encoder.py", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
out = Path(args.output_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Loading {CLIP_MODEL} ...")
|
||||
model = CLIPTextModelWithProjection.from_pretrained(CLIP_MODEL)
|
||||
tokenizer = CLIPTokenizer.from_pretrained(CLIP_MODEL)
|
||||
model.eval()
|
||||
|
||||
# --- Export ONNX text encoder ---
|
||||
# Input: input_ids (1, 77) int64 — BPE token IDs, padded/truncated to 77
|
||||
# Output: text_embeds (1, 512) float32 — projected features (NOT normalized; C++ normalizes)
|
||||
onnx_path = out / "clip_text_encoder.onnx"
|
||||
print(f"Exporting ONNX → {onnx_path} ...")
|
||||
dummy = torch.zeros(1, 77, dtype=torch.long)
|
||||
torch.onnx.export(
|
||||
model,
|
||||
dummy,
|
||||
str(onnx_path),
|
||||
input_names=["input_ids"],
|
||||
output_names=["text_embeds"],
|
||||
dynamic_axes={
|
||||
"input_ids": {0: "batch_size"},
|
||||
"text_embeds": {0: "batch_size"},
|
||||
},
|
||||
opset_version=14,
|
||||
)
|
||||
|
||||
# --- Save vocab as tab-separated text (token\tid) ---
|
||||
# Avoids JSON parsing in C++. Tokens with special chars are escaped.
|
||||
vocab_path = out / "clip_vocab.txt"
|
||||
print(f"Saving vocab → {vocab_path} ...")
|
||||
vocab = tokenizer.get_vocab()
|
||||
with open(vocab_path, "w", encoding="utf-8") as f:
|
||||
for token, idx in sorted(vocab.items(), key=lambda x: x[1]):
|
||||
token_esc = (token
|
||||
.replace("\\", "\\\\")
|
||||
.replace("\t", "\\t")
|
||||
.replace("\n", "\\n"))
|
||||
f.write(f"{token_esc}\t{idx}\n")
|
||||
|
||||
# --- Save BPE merge rules ---
|
||||
merges_path = out / "clip_merges.txt"
|
||||
print(f"Saving merges → {merges_path} ...")
|
||||
with open(merges_path, "w", encoding="utf-8") as f:
|
||||
for (a, b) in tokenizer.bpe_ranks.keys():
|
||||
f.write(f"{a} {b}\n")
|
||||
|
||||
print(f"\nDone. Files written to {out.resolve()}/")
|
||||
print(f" clip_text_encoder.onnx — ONNX model")
|
||||
print(f" clip_vocab.txt — {len(vocab)} tokens")
|
||||
print(f" clip_merges.txt — {len(tokenizer.bpe_ranks)} merge rules")
|
||||
print(f"\nSet env vars:")
|
||||
print(f" ZED_ONNX_TEXT_ENCODER_PATH={out.resolve()}/clip_text_encoder.onnx")
|
||||
print(f" ZED_ONNX_CLIP_VOCAB={out.resolve()}/clip_vocab.txt")
|
||||
print(f" ZED_ONNX_CLIP_MERGES={out.resolve()}/clip_merges.txt")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
Executable
+75
@@ -0,0 +1,75 @@
|
||||
#!/usr/bin/env bash
|
||||
# Install build dependencies for dora-zed on NVIDIA Jetson (ARM64 / L4T).
|
||||
# Skips any component that is already installed.
|
||||
# Uses aria2c for fast parallel downloads if available, otherwise falls back to wget.
|
||||
#
|
||||
# Supported JetPack versions for ZED SDK 5.x:
|
||||
# JetPack 6.0 GA (L4T 36.3) — Jetson Orin
|
||||
# JetPack 6.1/6.2 (L4T 36.4)
|
||||
# JetPack 6.2.2 (L4T 36.5)
|
||||
# JetPack 7.0 (L4T 38.2) — Jetson Thor
|
||||
# JetPack 7.1 (L4T 38.4)
|
||||
#
|
||||
# ONNX Runtime: Microsoft does not publish official GPU binaries for ARM64.
|
||||
# The native engine (ZED SDK built-in TensorRT) works without ONNX Runtime.
|
||||
# To use the open_vocab engine, build ONNX Runtime from source:
|
||||
# https://onnxruntime.ai/docs/build/eps.html#tensorrt
|
||||
set -euo pipefail
|
||||
|
||||
ZED_VERSION="5.2"
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
TMP_DIR="${SCRIPT_DIR}/../tmp"
|
||||
mkdir -p "$TMP_DIR"
|
||||
|
||||
download() {
|
||||
local url="$1" dest="$2"
|
||||
if command -v aria2c &>/dev/null; then
|
||||
aria2c -x 16 -s 16 -d "$(dirname "$dest")" -o "$(basename "$dest")" "$url"
|
||||
else
|
||||
wget -O "$dest" "$url"
|
||||
fi
|
||||
}
|
||||
|
||||
# --- Detect L4T version ---
|
||||
if [ ! -f /etc/nv_tegra_release ]; then
|
||||
echo "Error: /etc/nv_tegra_release not found. Is this a Jetson device?" >&2
|
||||
exit 1
|
||||
fi
|
||||
L4T_VERSION=$(head -1 /etc/nv_tegra_release | grep -oP 'R\K[0-9]+\.[0-9]+' | head -1)
|
||||
L4T_MAJOR=$(echo "$L4T_VERSION" | cut -d'.' -f1)
|
||||
echo ">>> Detected L4T ${L4T_VERSION} (major: ${L4T_MAJOR})"
|
||||
|
||||
case "$L4T_MAJOR" in
|
||||
36) L4T_PATH="l4t36" ;;
|
||||
38) L4T_PATH="l4t38" ;;
|
||||
35) L4T_PATH="l4t35" ;;
|
||||
*)
|
||||
echo "Warning: L4T ${L4T_MAJOR} not tested. Attempting l4t${L4T_MAJOR}." >&2
|
||||
L4T_PATH="l4t${L4T_MAJOR}"
|
||||
;;
|
||||
esac
|
||||
|
||||
# --- ZED SDK ---
|
||||
if [ -d /usr/local/zed ]; then
|
||||
echo ">>> ZED SDK already installed at /usr/local/zed — skipping."
|
||||
else
|
||||
echo ">>> Downloading ZED SDK ${ZED_VERSION} for Jetson (${L4T_PATH})..."
|
||||
download \
|
||||
"https://download.stereolabs.com/zedsdk/${ZED_VERSION}/${L4T_PATH}/jetsons" \
|
||||
"${TMP_DIR}/ZED_SDK.run"
|
||||
chmod +x "${TMP_DIR}/ZED_SDK.run"
|
||||
echo ">>> Installing ZED SDK (silent)..."
|
||||
sudo "${TMP_DIR}/ZED_SDK.run" -- silent skip_python skip_tools
|
||||
fi
|
||||
|
||||
# --- ONNX Runtime ---
|
||||
echo ""
|
||||
echo ">>> ONNX Runtime (ARM64 GPU): no official pre-built binary available."
|
||||
echo " The native ZED engine works without ONNX Runtime."
|
||||
echo " To use the open_vocab engine, build from source:"
|
||||
echo " https://onnxruntime.ai/docs/build/eps.html#tensorrt"
|
||||
echo ""
|
||||
|
||||
echo ">>> Done. You can now run: mkdir -p build && cd build && cmake .. && make install"
|
||||
echo ">>> Downloaded files are in ${TMP_DIR}"
|
||||
Executable
+56
@@ -0,0 +1,56 @@
|
||||
#!/usr/bin/env bash
|
||||
# Install build dependencies for dora-zed on Ubuntu 24.04 x86_64 (CUDA 12).
|
||||
# Skips any component that is already installed.
|
||||
# Uses aria2c for fast parallel downloads if available, otherwise falls back to wget.
|
||||
# Downloads are placed in ./tmp relative to this script for inspection.
|
||||
set -euo pipefail
|
||||
|
||||
ZED_VERSION="5.2"
|
||||
ORT_VERSION="1.24.4"
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
TMP_DIR="${SCRIPT_DIR}/../tmp"
|
||||
mkdir -p "$TMP_DIR"
|
||||
|
||||
download() {
|
||||
local url="$1" dest="$2"
|
||||
if command -v aria2c &>/dev/null; then
|
||||
aria2c -x 16 -s 16 -d "$(dirname "$dest")" -o "$(basename "$dest")" "$url"
|
||||
else
|
||||
wget -O "$dest" "$url"
|
||||
fi
|
||||
}
|
||||
|
||||
# --- ZED SDK ---
|
||||
if [ -d /usr/local/zed ]; then
|
||||
echo ">>> ZED SDK already installed at /usr/local/zed — skipping."
|
||||
else
|
||||
echo ">>> Downloading ZED SDK ${ZED_VERSION} for Ubuntu 24.04 (CUDA 12)..."
|
||||
download \
|
||||
"https://download.stereolabs.com/zedsdk/${ZED_VERSION}/cu12/ubuntu24" \
|
||||
"${TMP_DIR}/ZED_SDK.run"
|
||||
chmod +x "${TMP_DIR}/ZED_SDK.run"
|
||||
echo ">>> Installing ZED SDK (silent)..."
|
||||
sudo "${TMP_DIR}/ZED_SDK.run" -- silent skip_python skip_tools
|
||||
fi
|
||||
|
||||
# --- ONNX Runtime (C++ GPU, installs to /usr/local) ---
|
||||
if [ -f /usr/local/lib64/libonnxruntime.so ] && [ -d /usr/local/include/onnxruntime ]; then
|
||||
echo ">>> ONNX Runtime already installed — skipping."
|
||||
else
|
||||
echo ">>> Downloading ONNX Runtime ${ORT_VERSION} GPU..."
|
||||
download \
|
||||
"https://github.com/microsoft/onnxruntime/releases/download/v${ORT_VERSION}/onnxruntime-linux-x64-gpu-${ORT_VERSION}.tgz" \
|
||||
"${TMP_DIR}/onnxruntime-linux-x64-gpu-${ORT_VERSION}.tgz"
|
||||
echo ">>> Extracting..."
|
||||
tar -xzf "${TMP_DIR}/onnxruntime-linux-x64-gpu-${ORT_VERSION}.tgz" -C "${TMP_DIR}"
|
||||
echo ">>> Installing ONNX Runtime to /usr/local..."
|
||||
sudo mkdir -p /usr/local/include/onnxruntime
|
||||
sudo cp -r "${TMP_DIR}/onnxruntime-linux-x64-gpu-${ORT_VERSION}/include/"* /usr/local/include/onnxruntime/
|
||||
sudo mkdir -p /usr/local/lib64
|
||||
sudo cp -r "${TMP_DIR}/onnxruntime-linux-x64-gpu-${ORT_VERSION}/lib/"* /usr/local/lib64/
|
||||
sudo ldconfig
|
||||
fi
|
||||
|
||||
echo ">>> Done. You can now run: mkdir -p build && cd build && cmake .. && make install"
|
||||
echo ">>> Downloaded files are in ${TMP_DIR}"
|
||||
Executable
+47
@@ -0,0 +1,47 @@
|
||||
# Install build dependencies for dora-zed on Windows x64.
|
||||
# Skips any component that is already installed.
|
||||
# Run from an elevated PowerShell prompt.
|
||||
#
|
||||
# Requirements before running:
|
||||
# - CUDA Toolkit 12.x: https://developer.nvidia.com/cuda-downloads
|
||||
# - CMake 3.21+: https://cmake.org/download/
|
||||
# - Rust toolchain: https://rustup.rs
|
||||
|
||||
$ErrorActionPreference = "Stop"
|
||||
|
||||
$ZED_INSTALL_DIR = "C:\Program Files (x86)\ZED SDK"
|
||||
$ORT_VERSION = "1.24.4"
|
||||
$ORT_NAME = "onnxruntime-win-x64-gpu_cuda13-$ORT_VERSION"
|
||||
$ORT_URL = "https://github.com/microsoft/onnxruntime/releases/download/v$ORT_VERSION/$ORT_NAME.zip"
|
||||
$ORT_DEST = "C:\onnxruntime"
|
||||
|
||||
# --- ZED SDK ---
|
||||
if (Test-Path $ZED_INSTALL_DIR) {
|
||||
Write-Host ">>> ZED SDK already installed at '$ZED_INSTALL_DIR' — skipping."
|
||||
} else {
|
||||
Write-Host ">>> ZED SDK: download the installer from the release page and run it."
|
||||
Write-Host " https://www.stereolabs.com/developers/release"
|
||||
Write-Host " After install, ZED SDK will be at '$ZED_INSTALL_DIR'."
|
||||
Write-Host ""
|
||||
}
|
||||
|
||||
# --- ONNX Runtime ---
|
||||
if (Test-Path "$ORT_DEST\lib\onnxruntime.dll") {
|
||||
Write-Host ">>> ONNX Runtime already installed at '$ORT_DEST' — skipping."
|
||||
} else {
|
||||
Write-Host ">>> Downloading ONNX Runtime $ORT_VERSION GPU (CUDA 13)..."
|
||||
$tmp = "$env:TEMP\$ORT_NAME.zip"
|
||||
Invoke-WebRequest -Uri $ORT_URL -OutFile $tmp
|
||||
|
||||
Write-Host ">>> Extracting to $ORT_DEST..."
|
||||
Expand-Archive -Path $tmp -DestinationPath $env:TEMP -Force
|
||||
if (Test-Path $ORT_DEST) { Remove-Item $ORT_DEST -Recurse -Force }
|
||||
Move-Item "$env:TEMP\$ORT_NAME" $ORT_DEST
|
||||
Remove-Item $tmp
|
||||
|
||||
Write-Host ">>> ONNX Runtime installed to $ORT_DEST"
|
||||
Write-Host " Pass it to CMake: cmake .. -DCMAKE_PREFIX_PATH=C:\onnxruntime"
|
||||
}
|
||||
|
||||
Write-Host ""
|
||||
Write-Host ">>> Done."
|
||||
Reference in New Issue
Block a user