remove the YOLO-World open-vocabulary engine
Deletes the yoloworld_onnx engine, the CLIP BPE tokenizer, the standalone test binary and both export scripts. Drops the update_vocabulary input and the yoloworld_onnx/open_vocab dispatch from main.cc, along with the now unused Engine::set_labels. Documentation follows.
This commit is contained in:
@@ -1,105 +0,0 @@
|
||||
# /// 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()
|
||||
@@ -1,124 +0,0 @@
|
||||
# /// 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()
|
||||
@@ -12,7 +12,7 @@
|
||||
#
|
||||
# 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:
|
||||
# To use the yolo_onnx engine, build ONNX Runtime from source:
|
||||
# https://onnxruntime.ai/docs/build/eps.html#tensorrt
|
||||
set -euo pipefail
|
||||
|
||||
@@ -67,7 +67,7 @@ fi
|
||||
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 " To use the yolo_onnx engine, build from source:"
|
||||
echo " https://onnxruntime.ai/docs/build/eps.html#tensorrt"
|
||||
echo ""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user