cristhian aguilera
fd6993e294
- 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
106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
# /// 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()
|