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