From fd6993e2940594306f1de3228eee73a2581e51c6 Mon Sep 17 00:00:00 2001 From: cristhian aguilera Date: Sat, 28 Mar 2026 10:41:11 -0300 Subject: [PATCH] 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 --- .gitignore | 4 + AGENTS.md | 14 +- CMakeLists.txt | 34 +- README.md | 109 +++++- scripts/export_clip_text_encoder.py | 105 +++++ scripts/export_yoloworld.py | 124 ++++++ scripts/install_linux_arm.sh | 75 ++++ scripts/install_linux_x64.sh | 56 +++ scripts/install_windows.ps1 | 47 +++ src/engines/engine.h | 27 ++ src/engines/open_vocab.cc | 587 ++++++++++++++++++++++++++++ src/engines/open_vocab.h | 46 +++ src/engines/zed_native.cc | 4 + src/engines/zed_native.h | 21 +- src/main.cc | 27 +- src/test_open_vocab.cc | 142 +++++++ src/util/clip_tokenizer.h | 226 +++++++++++ test/assets/.gitkeep | 0 18 files changed, 1611 insertions(+), 37 deletions(-) create mode 100644 scripts/export_clip_text_encoder.py create mode 100644 scripts/export_yoloworld.py create mode 100755 scripts/install_linux_arm.sh create mode 100755 scripts/install_linux_x64.sh create mode 100755 scripts/install_windows.ps1 create mode 100644 src/engines/engine.h create mode 100644 src/engines/open_vocab.cc create mode 100644 src/engines/open_vocab.h create mode 100644 src/test_open_vocab.cc create mode 100644 src/util/clip_tokenizer.h create mode 100644 test/assets/.gitkeep diff --git a/.gitignore b/.gitignore index d5697ec..d3aa2de 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,6 @@ build/ bin/ +tmp/ +test/assets/*.png +test/assets/*.jpg +test/assets/*_result.* diff --git a/AGENTS.md b/AGENTS.md index b9d0158..086a9fd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -14,7 +14,7 @@ mkdir -p build && cd build && cmake .. && make install The binary is installed to `bin/dora_zed` relative to the repo root. -**Requirements:** CMake 3.21+, ZED SDK 4.x (at `/usr/local/zed`), CUDA Toolkit, OpenCV 4.x, Rust stable toolchain (for dora-rs C++ bindings). +**Requirements:** CMake 3.21+, ZED SDK 5.x (at `/usr/local/zed`), CUDA Toolkit, ONNX Runtime 1.20+ (system install), Rust stable toolchain (for dora-rs C++ bindings). The first build clones dora-rs from GitHub and compiles `dora-node-api-cxx` via Cargo. This takes a while. Subsequent builds are incremental. If `DORA_ROOT_DIR` is set to a local dora checkout, it builds against that instead. @@ -32,16 +32,22 @@ Camera and detection behavior is configured entirely through environment variabl - **`src/util/env.h`** — Shared inline env-variable parsing helpers (`get_env_int`, `get_env_str`, `get_env_bool`, `get_env_float`). - **`src/zed/camera.h`** — `zed::Config`, `zed::Calibration`, and `zed::Camera` class. Camera wraps the ZED SDK lifecycle: open (with retry loop), warmup, grab (image + point cloud), calibration, and close. Exposes `handle()` for engine access. - **`src/zed/camera.cc`** — Implementation. SDK enum parsers live in an anonymous namespace. -- **`src/engines/zed_native.h`** — `engines::DetectionConfig`, `engines::DetectedObject`, and `engines::ZedNative` class. Wraps ZED SDK's TensorRT-optimized object detection module. -- **`src/engines/zed_native.cc`** — Implementation. Enables detection on the camera, retrieves objects per frame, converts to `DetectedObject` structs. -- **`src/main.cc`** — Dora event loop. Receives `tick` inputs, calls `camera.grab()` and `detector.detect()`, then sends `image_bgr`, `point_cloud`, `camera_info`, and `objects` outputs with shared `frame_id`. Handles reconnection on grab failure. +- **`src/engines/engine.h`** — `engines::DetectedObject` struct and `engines::Engine` abstract base class (`enable`, `detect`, `objects`, `label_names`). +- **`src/engines/zed_native.h/cc`** — `engines::ZedNative`: wraps ZED SDK's TensorRT-optimized object detection module. +- **`src/engines/open_vocab.h/cc`** — `engines::OpenVocab`: YOLO-World open-vocabulary detection via ONNX Runtime. Supports two modes: single-input (fixed vocab, legacy) and two-input (open-vocab, `images` + `text_feats`). In two-input mode, reads `ZED_ONNX_LABELS` on every tick; when the labels change, re-tokenizes and re-encodes via the CLIP text encoder ONNX session, then passes the resulting `(1, N, 512)` L2-normalized embeddings to the visual model. Preprocesses images (letterbox, normalize, CHW), runs inference, applies NMS, then fuses 2D detections with the ZED point cloud to compute 3D positions and bounding boxes. +- **`src/util/clip_tokenizer.h`** — Header-only C++ CLIP BPE tokenizer (no external deps). Implements the GPT-2 `bytes_to_unicode()` byte encoder + iterative BPE merge algorithm. Loads `clip_vocab.txt` and `clip_merges.txt` generated by `scripts/export_clip_text_encoder.py`. Returns 77 `int64` token IDs (BOS=49406, EOS=49407, padded with 0). +- **`src/main.cc`** — Dora event loop. Selects engine via `ZED_ENGINE` env var, holds a `std::unique_ptr`. Receives `tick` inputs, calls `camera.grab()` and `detector->detect()`, then sends `image_bgr`, `point_cloud`, `camera_info`, and `objects` outputs with shared `frame_id`. Handles reconnection on grab failure. - **`DoraTargets.cmake`** — Handles building or fetching dora-rs and generating the CXX bridge files. Two paths: local (`DORA_ROOT_DIR`) or git clone. Produces `dora_node_api_cxx` library and the `node_bridge.cc` glue. - **`dataflow.yml`** — Example dora dataflow definition using the `git:` source format. +- **`scripts/export_yoloworld.py`** — PEP 723 uv script. Exports a YOLOv8-worldv2 model to ONNX with N class slots (labels arbitrary at runtime). Run: `uv run scripts/export_yoloworld.py --model s --num-classes 10`. +- **`scripts/export_clip_text_encoder.py`** — PEP 723 uv script. Exports CLIP ViT-B/32 text encoder to ONNX and saves `clip_vocab.txt` / `clip_merges.txt`. Run once; reused across all YOLO-World model sizes. Run: `uv run scripts/export_clip_text_encoder.py --output-dir models/`. +- **`src/test_open_vocab.cc`** — Standalone C++ test binary (`bin/test_open_vocab`). Exercises the `OpenVocab` engine with images from `test/assets/yolo.png` and `test/assets/yoloworld.png` (paths overridable via argv). No dora and no live camera required — `enable()` ignores the camera for this engine. Uses the same env vars as the main binary. Saves annotated result images alongside inputs. ## Key APIs - **dora C++ API** (``): `init_dora_node()`, `dora_node.events->next()`, `event_type()`, `send_output_with_metadata()`, `new_metadata()`. - **ZED SDK** (``): `sl::Camera`, `sl::InitParameters`, `sl::RuntimeParameters`, `sl::Mat`, `sl::ObjectDetectionParameters`, `sl::Objects`. +- **ONNX Runtime** (``): `Ort::Env`, `Ort::Session`, `Ort::SessionOptions`, `Ort::Value`. Used by the `open_vocab` engine. - Output data is sent as raw byte slices (`rust::Slice`) with typed metadata (shape, dtype, encoding). ## Tools diff --git a/CMakeLists.txt b/CMakeLists.txt index e59b311..8310795 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -13,10 +13,20 @@ find_package(zed REQUIRED) find_package(OpenCV REQUIRED) +find_package(onnxruntime REQUIRED) +set(ONNXRUNTIME_INCLUDE_DIRS "") # propagated via imported target +set(ONNXRUNTIME_LIBRARIES onnxruntime::onnxruntime) + link_directories(${dora_link_dirs}) link_directories(${ZED_LIBRARY_DIR}) -add_executable(dora_zed src/main.cc src/zed/camera.cc src/engines/zed_native.cc ${node_bridge}) +add_executable(dora_zed + src/main.cc + src/zed/camera.cc + src/engines/zed_native.cc + src/engines/open_vocab.cc + ${node_bridge} +) add_dependencies(dora_zed Dora_cxx) target_include_directories( @@ -28,8 +38,28 @@ target_include_directories( ${ZED_INCLUDE_DIRS} ${OpenCV_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS} + ${ONNXRUNTIME_INCLUDE_DIRS} ) -target_link_libraries(dora_zed dora_node_api_cxx ${ZED_LIBRARIES} ${OpenCV_LIBS}) +target_link_libraries(dora_zed dora_node_api_cxx ${ZED_LIBRARIES} ${OpenCV_LIBS} ${ONNXRUNTIME_LIBRARIES}) install(TARGETS dora_zed DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/bin) + +# Standalone detector test — no dora, no live camera required +add_executable(test_open_vocab + src/test_open_vocab.cc + src/engines/open_vocab.cc +) +target_include_directories(test_open_vocab PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${ZED_INCLUDE_DIRS} + ${OpenCV_INCLUDE_DIRS} + ${CUDAToolkit_INCLUDE_DIRS} + ${ONNXRUNTIME_INCLUDE_DIRS} +) +target_link_libraries(test_open_vocab PRIVATE + ${ZED_LIBRARIES} + ${OpenCV_LIBS} + ${ONNXRUNTIME_LIBRARIES} +) +install(TARGETS test_open_vocab DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/bin) diff --git a/README.md b/README.md index 49c99ea..4b84ffa 100644 --- a/README.md +++ b/README.md @@ -43,11 +43,39 @@ The binary is installed to `bin/dora_zed` relative to the repository root. | Dependency | Minimum version | Notes | |---|---|---| -| CMake | 3.21 | | -| ZED SDK | 4.x | Default path: `/usr/local/zed` | -| CUDA Toolkit | Required by ZED SDK | | -| OpenCV | 4.x | | -| Rust toolchain | stable | Used to build dora-rs C++ bindings | +| [CMake](https://cmake.org/download/) | 3.21 | | +| [ZED SDK](https://www.stereolabs.com/developers/release) | 5.x | Default path: `/usr/local/zed`; bundles OpenCV | +| [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) | Required by ZED SDK | | +| [Rust toolchain](https://rustup.rs) | stable | Used to build dora-rs C++ bindings | +| [ONNX Runtime](https://github.com/microsoft/onnxruntime/releases) | 1.20+ | Must be installed before building; required only for `open_vocab` engine | + +## Installation scripts + +Helper scripts are in [`scripts/`](scripts/): + +| Platform | Script | Notes | +|---|---|---| +| Ubuntu 24.04 x86_64 | [`scripts/install_linux_x64.sh`](scripts/install_linux_x64.sh) | Installs ZED SDK + ONNX Runtime | +| NVIDIA Jetson (ARM64) | [`scripts/install_linux_arm.sh`](scripts/install_linux_arm.sh) | Installs ZED SDK; auto-detects L4T version | +| Windows x64 | [`scripts/install_windows.ps1`](scripts/install_windows.ps1) | Installs ONNX Runtime; ZED SDK downloaded manually | + +**Linux / Jetson:** + +```bash +chmod +x scripts/install_linux_x64.sh # or install_linux_arm.sh +sudo scripts/install_linux_x64.sh +``` + +**Windows** (elevated PowerShell): + +```powershell +Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass +.\scripts\install_windows.ps1 +``` + +> **Jetson note:** No official ONNX Runtime GPU binary exists for ARM64. The `native` engine (ZED SDK built-in TensorRT) works without it. The `open_vocab` engine requires [building ONNX Runtime from source](https://onnxruntime.ai/docs/build/eps.html#tensorrt). + +> **Windows note:** ZED SDK must be downloaded and installed manually from [stereolabs.com/developers/release](https://www.stereolabs.com/developers/release). Pass the ONNX Runtime path to CMake: `cmake .. -DCMAKE_PREFIX_PATH=C:\onnxruntime` ## Environment variables @@ -62,22 +90,35 @@ The binary is installed to `bin/dora_zed` relative to the repository root. | `ZED_FLIP` | `ON` | Image flip: `ON` or `OFF` | | `ZED_WARMUP_FRAMES` | `30` | Number of warmup frames to discard | | `ZED_RETRY_INTERVAL_SECONDS` | `5` | Seconds between reconnection attempts | -| `ZED_DETECTION_MODEL` | `MULTI_CLASS_BOX_MEDIUM` | Detection model: `MULTI_CLASS_BOX_FAST`, `MULTI_CLASS_BOX_MEDIUM`, or `MULTI_CLASS_BOX_ACCURATE` | -| `ZED_DETECTION_CONFIDENCE` | `40` | Detection confidence threshold (0-100) | -| `ZED_DETECTION_TRACKING` | `false` | Enable object tracking across frames (enables positional tracking) | +| `ZED_ENGINE` | `native` | Detection engine: `native` (ZED SDK) or `open_vocab` (YOLO-World via ONNX Runtime) | +| `ZED_DETECTION_MODEL` | `MULTI_CLASS_BOX_MEDIUM` | (`native`) Detection model: `MULTI_CLASS_BOX_FAST`, `MULTI_CLASS_BOX_MEDIUM`, or `MULTI_CLASS_BOX_ACCURATE` | +| `ZED_DETECTION_CONFIDENCE` | `40` | (`native`) Detection confidence threshold (0-100) | +| `ZED_DETECTION_TRACKING` | `false` | (`native`) Enable object tracking across frames | +| `ZED_ONNX_MODEL_PATH` | _(required)_ | (`open_vocab`) Path to YOLO-World ONNX model file | +| `ZED_ONNX_LABELS` | _(required)_ | (`open_vocab`) Comma-separated label names (e.g. `person,cup,chair`). Can be changed between ticks — the node re-encodes embeddings automatically | +| `ZED_ONNX_CONFIDENCE` | `0.25` | (`open_vocab`) Detection confidence threshold (0.0-1.0) | +| `ZED_ONNX_NMS_THRESHOLD` | `0.45` | (`open_vocab`) NMS IoU threshold | +| `ZED_ONNX_INPUT_SIZE` | `640` | (`open_vocab`) Model input image size in pixels | +| `ZED_ONNX_TEXT_ENCODER_PATH` | _(required for open-vocab)_ | (`open_vocab`) Path to `clip_text_encoder.onnx`. Required when the visual model has two inputs (`images` + `text_feats`) | +| `ZED_ONNX_CLIP_VOCAB` | _(required for open-vocab)_ | (`open_vocab`) Path to `clip_vocab.txt` BPE vocabulary file | +| `ZED_ONNX_CLIP_MERGES` | _(required for open-vocab)_ | (`open_vocab`) Path to `clip_merges.txt` BPE merge rules file | ## Outputs ### `image_bgr` + BGR image as raw bytes. Metadata includes `shape` (HxWx3), `encoding` (`bgr8`), camera intrinsics (`fx`, `fy`, `cx`, `cy`), `distortion`, `frame_id`, and `timestamp_ns`. ### `point_cloud` + XYZRGBA point cloud as float32 bytes. Metadata includes `shape` (HxWx4), `units` (`mm`), and camera intrinsics. ### `camera_info` + Four float32 values: `[fx, fy, cx, cy]`. Metadata includes `distortion` coefficients. ### `objects` + Detected objects as a flat float32 array with shape `[num_objects, 37]`. Each row contains 37 floats: | Index | Field | Description | @@ -88,4 +129,54 @@ Detected objects as a flat float32 array with shape `[num_objects, 37]`. Each ro | 5-12 | `bb2d_*` | 2D bounding box: 4 corners × 2 pixel coords | | 13-36 | `bb3d_*` | 3D bounding box: 8 corners × 3 world coords (mm) | -Metadata includes `frame_id` and `timestamp_ns` matching the corresponding `image_bgr`, plus `image_width` and `image_height` for interpreting 2D bounding box pixel coordinates. When no objects are detected, an empty array with shape `[0, 37]` is sent. +Metadata includes `frame_id` and `timestamp_ns` matching the corresponding `image_bgr`, `image_width` and `image_height` for interpreting 2D coordinates, `engine` (`native` or `open_vocab`), and `label_names` (comma-separated class names for the active engine). When no objects are detected, an empty array with shape `[0, 37]` is sent. + +**`open_vocab` engine:** Uses a YOLO-World model exported with open-vocabulary support (two ONNX inputs: `images` and `text_feats`). Labels are set at runtime via `ZED_ONNX_LABELS` and re-encoded automatically on each tick when they change — no node restart required. Export the model and CLIP text encoder with the helper scripts: + +```bash +# Export the visual model (N=10 class slots, can be changed) +uv run scripts/export_yoloworld.py --model s --num-classes 10 --output-dir models/ + +# Export the CLIP text encoder + vocab files (once, reused for all model sizes) +uv run scripts/export_clip_text_encoder.py --output-dir models/ + +# Run +ZED_ENGINE=open_vocab \ +ZED_ONNX_MODEL_PATH=models/yolov8s-worldv2_n10_i640.onnx \ +ZED_ONNX_TEXT_ENCODER_PATH=models/clip_text_encoder.onnx \ +ZED_ONNX_CLIP_VOCAB=models/clip_vocab.txt \ +ZED_ONNX_CLIP_MERGES=models/clip_merges.txt \ +ZED_ONNX_LABELS="person,car,bag,chair,laptop" \ +dora start dataflow.yml +``` + +To verify the model works before connecting the camera, build and run the standalone C++ test binary (`bin/test_open_vocab`). It requires no live camera — images are loaded from disk: + +```bash +cd build && cmake .. && make test_open_vocab + +# Place test images: +cp street.jpg test/assets/yolo.png +cp custom.jpg test/assets/yoloworld.png + +# Fixed-vocab YOLO (both images use the same model): +ZED_ONNX_MODEL_PATH=models/yolov8s_n80_i640.onnx \ +ZED_ONNX_LABELS="person,car,bus,bicycle" \ +bin/test_open_vocab + +# Open-vocab YOLO-World (separate model for second image via ZED_ONNX_WORLD_MODEL_PATH): +ZED_ONNX_MODEL_PATH=models/yolov8s_n80_i640.onnx \ +ZED_ONNX_WORLD_MODEL_PATH=models/yolov8s-worldv2_n10_i640.onnx \ +ZED_ONNX_TEXT_ENCODER_PATH=models/clip_text_encoder.onnx \ +ZED_ONNX_CLIP_VOCAB=models/clip_vocab.txt \ +ZED_ONNX_CLIP_MERGES=models/clip_merges.txt \ +ZED_ONNX_LABELS="person,car,bus,bicycle" \ +bin/test_open_vocab + +# Override image paths: +bin/test_open_vocab /path/to/img1.jpg /path/to/img2.jpg +``` + +Annotated results are saved alongside the input images as `yolo_result.png` and `yoloworld_result.png`. + +3D positions are computed by fusing the 2D detections with the ZED point cloud (median of valid depth samples within each bounding box). 3D bounding box corners are back-projected at the p25/p75 depth percentiles. diff --git a/scripts/export_clip_text_encoder.py b/scripts/export_clip_text_encoder.py new file mode 100644 index 0000000..d6216cc --- /dev/null +++ b/scripts/export_clip_text_encoder.py @@ -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() diff --git a/scripts/export_yoloworld.py b/scripts/export_yoloworld.py new file mode 100644 index 0000000..df28594 --- /dev/null +++ b/scripts/export_yoloworld.py @@ -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() diff --git a/scripts/install_linux_arm.sh b/scripts/install_linux_arm.sh new file mode 100755 index 0000000..8faa713 --- /dev/null +++ b/scripts/install_linux_arm.sh @@ -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}" diff --git a/scripts/install_linux_x64.sh b/scripts/install_linux_x64.sh new file mode 100755 index 0000000..cd10f5d --- /dev/null +++ b/scripts/install_linux_x64.sh @@ -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}" diff --git a/scripts/install_windows.ps1 b/scripts/install_windows.ps1 new file mode 100755 index 0000000..8f8979e --- /dev/null +++ b/scripts/install_windows.ps1 @@ -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." diff --git a/src/engines/engine.h b/src/engines/engine.h new file mode 100644 index 0000000..eb3f07a --- /dev/null +++ b/src/engines/engine.h @@ -0,0 +1,27 @@ +#pragma once + +#include +#include +#include + +namespace engines { + +struct DetectedObject { + int label_id; + float confidence; + float position[3]; // center x, y, z (mm) + float bounding_box_2d[8]; // 4 corners × 2 coords (pixels) + float bounding_box_3d[24]; // 8 corners × 3 coords (mm) +}; + +class Engine { +public: + virtual ~Engine() = default; + + virtual void enable(sl::Camera& camera) = 0; + virtual void detect(sl::Camera& camera) = 0; + virtual const std::vector& objects() const = 0; + virtual std::string label_names() const = 0; +}; + +} // namespace engines diff --git a/src/engines/open_vocab.cc b/src/engines/open_vocab.cc new file mode 100644 index 0000000..b1851fd --- /dev/null +++ b/src/engines/open_vocab.cc @@ -0,0 +1,587 @@ +#include "engines/open_vocab.h" +#include "util/clip_tokenizer.h" +#include "util/env.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// --- NMS helpers --- + +struct RawDetection { + float x1, y1, x2, y2; + float score; + int class_id; +}; + +float iou(const RawDetection& a, const RawDetection& b) { + float inter_x1 = std::max(a.x1, b.x1); + float inter_y1 = std::max(a.y1, b.y1); + float inter_x2 = std::min(a.x2, b.x2); + float inter_y2 = std::min(a.y2, b.y2); + float inter_w = std::max(0.0f, inter_x2 - inter_x1); + float inter_h = std::max(0.0f, inter_y2 - inter_y1); + float inter_area = inter_w * inter_h; + float area_a = (a.x2 - a.x1) * (a.y2 - a.y1); + float area_b = (b.x2 - b.x1) * (b.y2 - b.y1); + float union_area = area_a + area_b - inter_area; + return union_area > 0.0f ? inter_area / union_area : 0.0f; +} + +std::vector nms(std::vector dets, float threshold) { + std::sort(dets.begin(), dets.end(), + [](const RawDetection& a, const RawDetection& b) { + return a.score > b.score; + }); + std::vector suppressed(dets.size(), false); + std::vector result; + for (size_t i = 0; i < dets.size(); i++) { + if (suppressed[i]) continue; + result.push_back(dets[i]); + for (size_t j = i + 1; j < dets.size(); j++) { + if (!suppressed[j] && iou(dets[i], dets[j]) > threshold) { + suppressed[j] = true; + } + } + } + return result; +} + +// --- Letterbox resize --- +cv::Mat letterbox(const cv::Mat& src, int size, float& scale, float& pad_x, float& pad_y) { + float scale_w = static_cast(size) / src.cols; + float scale_h = static_cast(size) / src.rows; + scale = std::min(scale_w, scale_h); + int new_w = static_cast(std::round(src.cols * scale)); + int new_h = static_cast(std::round(src.rows * scale)); + pad_x = (size - new_w) / 2.0f; + pad_y = (size - new_h) / 2.0f; + + cv::Mat resized; + cv::resize(src, resized, {new_w, new_h}); + + cv::Mat canvas(size, size, CV_8UC3, cv::Scalar(114, 114, 114)); + resized.copyTo(canvas(cv::Rect( + static_cast(pad_x), static_cast(pad_y), new_w, new_h))); + return canvas; +} + +// --- Point cloud fusion --- +struct FusionResult { + float cx, cy, cz; + float depth_near, depth_far; + bool valid = false; +}; + +FusionResult fuse_bbox(const sl::Mat& pc, int x1, int y1, int x2, int y2) { + FusionResult result{}; + int w = static_cast(pc.getWidth()); + int h = static_cast(pc.getHeight()); + x1 = std::max(x1, 0); + y1 = std::max(y1, 0); + x2 = std::min(x2, w - 1); + y2 = std::min(y2, h - 1); + + std::vector xs, ys, zs; + xs.reserve((x2 - x1) * (y2 - y1)); + ys.reserve(xs.capacity()); + zs.reserve(xs.capacity()); + + for (int y = y1; y <= y2; y++) { + for (int x = x1; x <= x2; x++) { + sl::float4 pt; + pc.getValue(x, y, &pt); + if (!std::isnan(pt.x) && !std::isnan(pt.y) && !std::isnan(pt.z) && + !std::isinf(pt.x) && !std::isinf(pt.y) && !std::isinf(pt.z)) { + xs.push_back(pt.x); + ys.push_back(pt.y); + zs.push_back(pt.z); + } + } + } + + if (xs.empty()) return result; + + std::sort(xs.begin(), xs.end()); + std::sort(ys.begin(), ys.end()); + std::sort(zs.begin(), zs.end()); + + size_t n = xs.size(); + size_t mid = n / 2; + size_t p25 = n / 4; + size_t p75 = 3 * n / 4; + + result.cx = xs[mid]; + result.cy = ys[mid]; + result.cz = zs[mid]; + result.depth_near = zs[p25]; + result.depth_far = zs[p75]; + result.valid = true; + return result; +} + +void backproject(float u, float v, float z, float fx, float fy, float cx, float cy, + float& X, float& Y, float& Z) { + Z = z; + X = (u - cx) * z / fx; + Y = (v - cy) * z / fy; +} + +// Split a comma-separated string into a vector of trimmed tokens +std::vector split_labels(const std::string& s) { + std::vector out; + std::stringstream ss(s); + std::string tok; + while (std::getline(ss, tok, ',')) { + size_t l = tok.find_first_not_of(" \t"); + size_t r = tok.find_last_not_of(" \t"); + if (l != std::string::npos) + out.push_back(tok.substr(l, r - l + 1)); + } + return out; +} + +} // namespace + +namespace engines { + +// --- Impl --- + +struct OpenVocab::Impl { + // Visual model + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "open_vocab"}; + Ort::SessionOptions session_options; + std::unique_ptr session; + Ort::AllocatorWithDefaultOptions allocator; + + std::string input_name; // visual model: image input + std::string output_name; // visual model: detection output + std::string text_input_name; // visual model: text_feats input (empty if single-input model) + int num_classes = 0; + bool has_text_input = false; + + // Text encoder (optional; used when has_text_input == true) + Ort::SessionOptions text_session_options; + std::unique_ptr text_session; + std::string text_enc_input_name; // "input_ids" + std::string text_enc_output_name; // "text_embeds" + + // Tokenizer + embedding cache + util::ClipTokenizer tokenizer; + std::vector cached_text_feats; // (num_classes * 512) flat, row-major + std::string cached_labels; + + // Preprocessing + ONNX inference + NMS. + // Takes a BGR image (cv::imread format); returns 2D detections in original image coords. + std::vector run_inference(const OpenVocabConfig& config, + const cv::Mat& image_bgr); +}; + +// --- Impl::run_inference --- + +std::vector OpenVocab::Impl::run_inference( + const OpenVocabConfig& config, const cv::Mat& image_bgr) +{ + if (!session) return {}; + if (has_text_input && cached_text_feats.empty()) return {}; + + int orig_w = image_bgr.cols; + int orig_h = image_bgr.rows; + + // BGR → RGB (YOLO expects RGB) + cv::Mat rgb; + cv::cvtColor(image_bgr, rgb, cv::COLOR_BGR2RGB); + + float scale, pad_x, pad_y; + cv::Mat lb = letterbox(rgb, config.input_size, scale, pad_x, pad_y); + + cv::Mat lb_float; + lb.convertTo(lb_float, CV_32F, 1.0f / 255.0f); + + int sz = config.input_size; + std::vector input_data(3 * sz * sz); + for (int c = 0; c < 3; c++) { + for (int h = 0; h < sz; h++) { + for (int w = 0; w < sz; w++) { + input_data[c * sz * sz + h * sz + w] = + lb_float.at(h, w)[c]; + } + } + } + + auto mem_info = Ort::MemoryInfo::CreateCpu( + OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault); + + std::vector image_shape{1, 3, sz, sz}; + Ort::Value image_tensor = Ort::Value::CreateTensor( + mem_info, input_data.data(), input_data.size(), + image_shape.data(), image_shape.size()); + + const char* output_name_cstr = output_name.c_str(); + std::vector outputs; + + try { + if (has_text_input) { + std::vector text_shape{1, num_classes, 512}; + Ort::Value text_tensor = Ort::Value::CreateTensor( + mem_info, + cached_text_feats.data(), + cached_text_feats.size(), + text_shape.data(), text_shape.size()); + + const char* input_names[2] = { + input_name.c_str(), + text_input_name.c_str() + }; + Ort::Value input_tensors[2] = { + std::move(image_tensor), + std::move(text_tensor) + }; + outputs = session->Run( + Ort::RunOptions{nullptr}, + input_names, input_tensors, 2, + &output_name_cstr, 1); + } else { + const char* input_name_cstr = input_name.c_str(); + outputs = session->Run( + Ort::RunOptions{nullptr}, + &input_name_cstr, &image_tensor, 1, + &output_name_cstr, 1); + } + } catch (const Ort::Exception& e) { + std::cerr << "OpenVocab: inference failed: " << e.what() << std::endl; + return {}; + } + + // Post-process: output shape (1, 4+num_classes, num_anchors) + auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape(); + if (out_shape.size() < 3) return {}; + + int num_anchors = static_cast(out_shape[2]); + int nc = num_classes; + if (nc <= 0) return {}; + + const float* out = outputs[0].GetTensorData(); + + std::vector candidates; + for (int a = 0; a < num_anchors; a++) { + float cx_lb = out[0 * num_anchors + a]; + float cy_lb = out[1 * num_anchors + a]; + float w_lb = out[2 * num_anchors + a]; + float h_lb = out[3 * num_anchors + a]; + + float best_score = 0.0f; + int best_class = 0; + for (int c = 0; c < nc; c++) { + float s = out[(4 + c) * num_anchors + a]; + if (s > best_score) { best_score = s; best_class = c; } + } + if (best_score < config.confidence_threshold) continue; + + float x1_lb = cx_lb - w_lb / 2.0f; + float y1_lb = cy_lb - h_lb / 2.0f; + float x2_lb = cx_lb + w_lb / 2.0f; + float y2_lb = cy_lb + h_lb / 2.0f; + + float x1 = std::max(0.0f, std::min((x1_lb - pad_x) / scale, static_cast(orig_w - 1))); + float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast(orig_h - 1))); + float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast(orig_w - 1))); + float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast(orig_h - 1))); + + candidates.push_back({x1, y1, x2, y2, best_score, best_class}); + } + + return nms(std::move(candidates), config.nms_threshold); +} + +// --- Config --- + +OpenVocabConfig OpenVocabConfig::from_env() { + OpenVocabConfig c; + c.model_path = util::get_env_str("ZED_ONNX_MODEL_PATH", c.model_path); + c.confidence_threshold = util::get_env_float("ZED_ONNX_CONFIDENCE", c.confidence_threshold); + c.nms_threshold = util::get_env_float("ZED_ONNX_NMS_THRESHOLD", c.nms_threshold); + c.input_size = util::get_env_int("ZED_ONNX_INPUT_SIZE", c.input_size); + c.labels = util::get_env_str("ZED_ONNX_LABELS", c.labels); + c.text_encoder_path = util::get_env_str("ZED_ONNX_TEXT_ENCODER_PATH", c.text_encoder_path); + c.clip_vocab_path = util::get_env_str("ZED_ONNX_CLIP_VOCAB", c.clip_vocab_path); + c.clip_merges_path = util::get_env_str("ZED_ONNX_CLIP_MERGES", c.clip_merges_path); + return c; +} + +// --- OpenVocab --- + +OpenVocab::OpenVocab(const OpenVocabConfig& config) + : config_(config), impl_(std::make_unique()) {} + +OpenVocab::~OpenVocab() = default; + +void OpenVocab::enable(sl::Camera& /* camera */) { + if (config_.model_path.empty()) { + std::cerr << "OpenVocab: ZED_ONNX_MODEL_PATH is not set." << std::endl; + return; + } + + // --- Load visual model --- + OrtCUDAProviderOptions cuda_opts{}; + cuda_opts.device_id = 0; + try { + impl_->session_options.AppendExecutionProvider_CUDA(cuda_opts); + } catch (const Ort::Exception& e) { + std::cerr << "OpenVocab: CUDA unavailable (" << e.what() << "), using CPU." << std::endl; + } + impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); + + std::cout << "OpenVocab: loading visual model (" << config_.model_path << ")..." << std::endl; + try { + impl_->session = std::make_unique( + impl_->env, config_.model_path.c_str(), impl_->session_options); + } catch (const Ort::Exception& e) { + std::cerr << "OpenVocab: failed to load visual model: " << e.what() << std::endl; + return; + } + + impl_->input_name = impl_->session->GetInputNameAllocated(0, impl_->allocator).get(); + impl_->output_name = impl_->session->GetOutputNameAllocated(0, impl_->allocator).get(); + + size_t num_inputs = impl_->session->GetInputCount(); + impl_->has_text_input = (num_inputs >= 2); + if (impl_->has_text_input) { + impl_->text_input_name = impl_->session->GetInputNameAllocated(1, impl_->allocator).get(); + } + + auto out_info = impl_->session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo(); + auto shape = out_info.GetShape(); + if (shape.size() == 3 && shape[1] > 4) { + impl_->num_classes = static_cast(shape[1]) - 4; + } + + std::cout << "OpenVocab: visual model loaded." + << " Inputs=" << num_inputs + << " Classes=" << impl_->num_classes + << (impl_->has_text_input ? " (open-vocab)" : " (fixed-vocab)") + << std::endl; + + // --- Load text encoder + tokenizer (only needed for open-vocab model) --- + if (impl_->has_text_input) { + if (config_.text_encoder_path.empty()) { + std::cerr << "OpenVocab: ZED_ONNX_TEXT_ENCODER_PATH is not set " + "(required for open-vocab model)." << std::endl; + return; + } + + impl_->text_session_options.SetGraphOptimizationLevel( + GraphOptimizationLevel::ORT_ENABLE_ALL); + + std::cout << "OpenVocab: loading text encoder (" << config_.text_encoder_path + << ")..." << std::endl; + try { + impl_->text_session = std::make_unique( + impl_->env, config_.text_encoder_path.c_str(), impl_->text_session_options); + } catch (const Ort::Exception& e) { + std::cerr << "OpenVocab: failed to load text encoder: " << e.what() << std::endl; + return; + } + + impl_->text_enc_input_name = + impl_->text_session->GetInputNameAllocated(0, impl_->allocator).get(); + impl_->text_enc_output_name = + impl_->text_session->GetOutputNameAllocated(0, impl_->allocator).get(); + + if (!config_.clip_vocab_path.empty() && !config_.clip_merges_path.empty()) { + try { + impl_->tokenizer.load(config_.clip_vocab_path, config_.clip_merges_path); + std::cout << "OpenVocab: tokenizer loaded." << std::endl; + } catch (const std::exception& e) { + std::cerr << "OpenVocab: failed to load tokenizer: " << e.what() << std::endl; + return; + } + } else { + std::cerr << "OpenVocab: ZED_ONNX_CLIP_VOCAB / ZED_ONNX_CLIP_MERGES not set." << std::endl; + return; + } + + if (!config_.labels.empty()) { + encode_labels(config_.labels); + } + } +} + +// Encode a comma-separated label string into cached_text_feats. +// Always fills exactly num_classes * 512 floats (pads with zeros if fewer labels). +void OpenVocab::encode_labels(const std::string& labels_str) { + if (!impl_->text_session || !impl_->tokenizer.is_loaded()) return; + + auto label_vec = split_labels(labels_str); + int N = impl_->num_classes; + impl_->cached_text_feats.assign(static_cast(N) * 512, 0.0f); + + auto mem_info = Ort::MemoryInfo::CreateCpu( + OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault); + + int encode_n = std::min(static_cast(label_vec.size()), N); + for (int i = 0; i < encode_n; i++) { + auto tokens = impl_->tokenizer.tokenize(label_vec[i]); + + std::vector input_shape{1, util::ClipTokenizer::SEQ_LEN}; + Ort::Value token_tensor = Ort::Value::CreateTensor( + mem_info, tokens.data(), tokens.size(), + input_shape.data(), input_shape.size()); + + const char* in = impl_->text_enc_input_name.c_str(); + const char* out = impl_->text_enc_output_name.c_str(); + std::vector outputs; + try { + outputs = impl_->text_session->Run( + Ort::RunOptions{nullptr}, &in, &token_tensor, 1, &out, 1); + } catch (const Ort::Exception& e) { + std::cerr << "OpenVocab: text encoder failed for \"" << label_vec[i] + << "\": " << e.what() << std::endl; + continue; + } + + const float* emb = outputs[0].GetTensorData(); + + // L2 normalize + float norm = 0.0f; + for (int j = 0; j < 512; j++) norm += emb[j] * emb[j]; + norm = std::sqrt(norm); + float inv_norm = (norm > 1e-8f) ? 1.0f / norm : 0.0f; + for (int j = 0; j < 512; j++) + impl_->cached_text_feats[i * 512 + j] = emb[j] * inv_norm; + } + + impl_->cached_labels = labels_str; + std::cout << "OpenVocab: encoded " << encode_n << " label(s): " << labels_str << std::endl; +} + +// --- detect(sl::Camera&) — production path with 3D fusion --- + +void OpenVocab::detect(sl::Camera& camera) { + detected_.clear(); + if (!impl_->session) return; + + // Re-encode if labels changed (on-the-fly vocabulary switching) + if (impl_->has_text_input) { + const char* env_labels = std::getenv("ZED_ONNX_LABELS"); + std::string current_labels = env_labels ? env_labels : config_.labels; + if (current_labels != impl_->cached_labels) { + encode_labels(current_labels); + } + if (impl_->cached_text_feats.empty()) return; + } + + // Retrieve image and point cloud from camera + sl::Mat img_mat; + camera.retrieveImage(img_mat, sl::VIEW::LEFT); + + sl::Mat pc_mat; + camera.retrieveMeasure(pc_mat, sl::MEASURE::XYZRGBA); + + sl::CameraInformation info = camera.getCameraInformation(); + auto calib = info.camera_configuration.calibration_parameters.left_cam; + float fx = calib.fx, fy = calib.fy, cx = calib.cx, cy = calib.cy; + + int orig_w = static_cast(img_mat.getWidth()); + int orig_h = static_cast(img_mat.getHeight()); + + // RGBA → BGR for run_inference + cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr()); + cv::Mat bgr; + cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR); + + auto kept = impl_->run_inference(config_, bgr); + + // 3D fusion + for (const auto& det : kept) { + FusionResult fusion = fuse_bbox(pc_mat, + static_cast(det.x1), static_cast(det.y1), + static_cast(det.x2), static_cast(det.y2)); + + DetectedObject obj{}; + obj.label_id = det.class_id; + obj.confidence = det.score * 100.0f; + + if (fusion.valid) { + obj.position[0] = fusion.cx; + obj.position[1] = fusion.cy; + obj.position[2] = fusion.cz; + } + + obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1; + obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1; + obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2; + obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2; + + if (fusion.valid) { + float corners_u[4] = {det.x1, det.x2, det.x2, det.x1}; + float corners_v[4] = {det.y1, det.y1, det.y2, det.y2}; + float depths[2] = {fusion.depth_near, fusion.depth_far}; + for (int face = 0; face < 2; face++) { + for (int c = 0; c < 4; c++) { + float X, Y, Z; + backproject(corners_u[c], corners_v[c], depths[face], + fx, fy, cx, cy, X, Y, Z); + int idx = (face * 4 + c) * 3; + obj.bounding_box_3d[idx + 0] = X; + obj.bounding_box_3d[idx + 1] = Y; + obj.bounding_box_3d[idx + 2] = Z; + } + } + } + + detected_.push_back(obj); + } +} + +// --- detect(const cv::Mat&) — image-only path, no 3D fusion --- + +void OpenVocab::detect(const cv::Mat& image_bgr) { + detected_.clear(); + if (!impl_->session) return; + + // Re-encode if labels changed + if (impl_->has_text_input) { + const char* env_labels = std::getenv("ZED_ONNX_LABELS"); + std::string current_labels = env_labels ? env_labels : config_.labels; + if (current_labels != impl_->cached_labels) { + encode_labels(current_labels); + } + if (impl_->cached_text_feats.empty()) return; + } + + auto kept = impl_->run_inference(config_, image_bgr); + + for (const auto& det : kept) { + DetectedObject obj{}; + obj.label_id = det.class_id; + obj.confidence = det.score * 100.0f; + // position stays zero (no point cloud available) + obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1; + obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1; + obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2; + obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2; + // bounding_box_3d stays zero + detected_.push_back(obj); + } +} + +const std::vector& OpenVocab::objects() const { + return detected_; +} + +std::string OpenVocab::label_names() const { + return impl_->cached_labels.empty() ? config_.labels : impl_->cached_labels; +} + +} // namespace engines diff --git a/src/engines/open_vocab.h b/src/engines/open_vocab.h new file mode 100644 index 0000000..8e586e9 --- /dev/null +++ b/src/engines/open_vocab.h @@ -0,0 +1,46 @@ +#pragma once + +#include "engines/engine.h" + +#include +#include +#include + +namespace engines { + +struct OpenVocabConfig { + std::string model_path; // ZED_ONNX_MODEL_PATH (required) + float confidence_threshold = 0.25f; // ZED_ONNX_CONFIDENCE + float nms_threshold = 0.45f; // ZED_ONNX_NMS_THRESHOLD + int input_size = 640; // ZED_ONNX_INPUT_SIZE + std::string labels; // ZED_ONNX_LABELS (comma-separated) + + // Open-vocabulary text encoder (required when model has two inputs) + std::string text_encoder_path; // ZED_ONNX_TEXT_ENCODER_PATH + std::string clip_vocab_path; // ZED_ONNX_CLIP_VOCAB + std::string clip_merges_path; // ZED_ONNX_CLIP_MERGES + + static OpenVocabConfig from_env(); +}; + +class OpenVocab : public Engine { +public: + explicit OpenVocab(const OpenVocabConfig& config); + ~OpenVocab() override; + + void enable(sl::Camera& camera) override; + void detect(sl::Camera& camera) override; + void detect(const cv::Mat& image_bgr); // image-only; no 3D fusion + const std::vector& objects() const override; + std::string label_names() const override; + +private: + OpenVocabConfig config_; + struct Impl; + std::unique_ptr impl_; + std::vector detected_; + + void encode_labels(const std::string& labels_str); +}; + +} // namespace engines diff --git a/src/engines/zed_native.cc b/src/engines/zed_native.cc index d475e01..efe6600 100644 --- a/src/engines/zed_native.cc +++ b/src/engines/zed_native.cc @@ -92,4 +92,8 @@ const std::vector& ZedNative::objects() const { return detected_; } +std::string ZedNative::label_names() const { + return "PERSON,VEHICLE,BAG,ANIMAL,ELECTRONICS,FRUIT_VEGETABLE,SPORT"; +} + } // namespace engines diff --git a/src/engines/zed_native.h b/src/engines/zed_native.h index 92db3dd..45df940 100644 --- a/src/engines/zed_native.h +++ b/src/engines/zed_native.h @@ -1,8 +1,8 @@ #pragma once -#include +#include "engines/engine.h" + #include -#include namespace engines { @@ -14,21 +14,14 @@ struct DetectionConfig { static DetectionConfig from_env(); }; -struct DetectedObject { - int label_id; - float confidence; - float position[3]; // center x, y, z (mm) - float bounding_box_2d[8]; // 4 corners × 2 coords (pixels) - float bounding_box_3d[24]; // 8 corners × 3 coords (mm) -}; - -class ZedNative { +class ZedNative : public Engine { public: explicit ZedNative(const DetectionConfig& config); - void enable(sl::Camera& camera); - void detect(sl::Camera& camera); - const std::vector& objects() const; + void enable(sl::Camera& camera) override; + void detect(sl::Camera& camera) override; + const std::vector& objects() const override; + std::string label_names() const override; private: DetectionConfig config_; diff --git a/src/main.cc b/src/main.cc index 50e1dbc..2668927 100644 --- a/src/main.cc +++ b/src/main.cc @@ -1,4 +1,7 @@ +#include "engines/engine.h" +#include "engines/open_vocab.h" #include "engines/zed_native.h" +#include "util/env.h" #include "zed/camera.h" #include @@ -8,6 +11,7 @@ #include #include #include +#include #include #include @@ -66,9 +70,16 @@ int main() { camera.warmup(); auto calib = camera.calibration(); - auto det_config = engines::DetectionConfig::from_env(); - engines::ZedNative detector(det_config); - detector.enable(camera.handle()); + std::string engine_name = util::get_env_str("ZED_ENGINE", "native"); + std::unique_ptr detector; + if (engine_name == "open_vocab") { + auto cfg = engines::OpenVocabConfig::from_env(); + detector = std::make_unique(cfg); + } else { + auto cfg = engines::DetectionConfig::from_env(); + detector = std::make_unique(cfg); + } + detector->enable(camera.handle()); int64_t frame_id = 0; while (true) { @@ -91,11 +102,11 @@ int main() { camera.close(); camera.open(); calib = camera.calibration(); - detector.enable(camera.handle()); + detector->enable(camera.handle()); continue; } - detector.detect(camera.handle()); + detector->detect(camera.handle()); sl::Mat& image_mat = camera.image(); sl::Mat& point_cloud_mat = camera.point_cloud(); @@ -181,7 +192,7 @@ int main() { } // --- objects --- - const auto& objects = detector.objects(); + const auto& objects = detector->objects(); auto obj_buffer = serialize_objects(objects); int num_objects = static_cast(objects.size()); @@ -201,8 +212,8 @@ int main() { "bb3d_x2,bb3d_y2,bb3d_z2,bb3d_x3,bb3d_y3,bb3d_z3," "bb3d_x4,bb3d_y4,bb3d_z4,bb3d_x5,bb3d_y5,bb3d_z5," "bb3d_x6,bb3d_y6,bb3d_z6,bb3d_x7,bb3d_y7,bb3d_z7"); - obj_metadata->set_string("label_names", - "PERSON,VEHICLE,BAG,ANIMAL,ELECTRONICS,FRUIT_VEGETABLE,SPORT"); + obj_metadata->set_string("label_names", detector->label_names()); + obj_metadata->set_string("engine", engine_name); obj_metadata->set_int("frame_id", current_frame_id); obj_metadata->set_int("timestamp_ns", timestamp_ns); obj_metadata->set_int("image_width", width); diff --git a/src/test_open_vocab.cc b/src/test_open_vocab.cc new file mode 100644 index 0000000..acbe3a8 --- /dev/null +++ b/src/test_open_vocab.cc @@ -0,0 +1,142 @@ +// Standalone test for the OpenVocab detection engine. +// No dora framework, no live ZED camera — loads images from disk. +// +// Usage: +// bin/test_open_vocab [yolo_image] [yoloworld_image] +// yolo_image default: test/assets/yolo.png +// yoloworld_image default: test/assets/yoloworld.png +// +// Uses the same env vars as the main binary: +// ZED_ONNX_MODEL_PATH path to YOLO ONNX model (required) +// ZED_ONNX_LABELS comma-separated labels (required) +// ZED_ONNX_CONFIDENCE detection threshold (default 0.25) +// ZED_ONNX_NMS_THRESHOLD NMS IoU threshold (default 0.45) +// ZED_ONNX_INPUT_SIZE model input size (default 640) +// ZED_ONNX_TEXT_ENCODER_PATH clip text encoder ONNX (open-vocab only) +// ZED_ONNX_CLIP_VOCAB clip_vocab.txt path (open-vocab only) +// ZED_ONNX_CLIP_MERGES clip_merges.txt path (open-vocab only) +// ZED_ONNX_WORLD_MODEL_PATH alternative ONNX model for the second image; +// falls back to ZED_ONNX_MODEL_PATH if not set + +#include "engines/open_vocab.h" +#include "util/env.h" + +#include +#include + +#include +#include +#include +#include +#include + +namespace fs = std::filesystem; + +namespace { + +std::vector parse_csv(const std::string& s) { + std::vector out; + std::istringstream ss(s); + std::string tok; + while (std::getline(ss, tok, ',')) { + size_t l = tok.find_first_not_of(" \t"); + size_t r = tok.find_last_not_of(" \t"); + if (l != std::string::npos) + out.push_back(tok.substr(l, r - l + 1)); + } + return out; +} + +void draw_detections(cv::Mat& img, + const std::vector& objects, + const std::vector& labels) { + for (const auto& obj : objects) { + int x1 = static_cast(obj.bounding_box_2d[0]); + int y1 = static_cast(obj.bounding_box_2d[1]); + int x2 = static_cast(obj.bounding_box_2d[4]); + int y2 = static_cast(obj.bounding_box_2d[5]); + + std::string name = (obj.label_id >= 0 && obj.label_id < static_cast(labels.size())) + ? labels[obj.label_id] : std::to_string(obj.label_id); + + char buf[64]; + std::snprintf(buf, sizeof(buf), " %.2f", obj.confidence / 100.0f); + std::string text = name + buf; + + cv::rectangle(img, {x1, y1}, {x2, y2}, {0, 200, 0}, 2); + cv::putText(img, text, {x1, std::max(y1 - 6, 14)}, + cv::FONT_HERSHEY_SIMPLEX, 0.55, {0, 200, 0}, 2, cv::LINE_AA); + } +} + +bool run_test(engines::OpenVocab& engine, const std::string& image_path) { + cv::Mat img = cv::imread(image_path); + if (img.empty()) { + std::cerr << "Error: cannot read image: " << image_path << std::endl; + return false; + } + + engine.detect(img); + const auto& objects = engine.objects(); + auto labels = parse_csv(engine.label_names()); + + std::cout << "[" << image_path << "] " << objects.size() << " detection(s):" << std::endl; + for (const auto& obj : objects) { + std::string name = (obj.label_id >= 0 && obj.label_id < static_cast(labels.size())) + ? labels[obj.label_id] : std::to_string(obj.label_id); + std::cout << " " << name + << " conf=" << obj.confidence / 100.0f + << " bbox=[" + << obj.bounding_box_2d[0] << "," << obj.bounding_box_2d[1] + << " -> " + << obj.bounding_box_2d[4] << "," << obj.bounding_box_2d[5] + << "]" << std::endl; + } + + cv::Mat annotated = img.clone(); + draw_detections(annotated, objects, labels); + + fs::path stem = fs::path(image_path).stem(); + fs::path parent = fs::path(image_path).parent_path(); + std::string result_path = (parent / (stem.string() + "_result.png")).string(); + cv::imwrite(result_path, annotated); + std::cout << "Saved: " << result_path << std::endl; + return true; +} + +} // namespace + +int main(int argc, char* argv[]) { + std::string yolo_path = (argc > 1) ? argv[1] : "test/assets/yolo.png"; + std::string yoloworld_path = (argc > 2) ? argv[2] : "test/assets/yoloworld.png"; + + auto config = engines::OpenVocabConfig::from_env(); + + // enable() ignores this parameter for OpenVocab + sl::Camera dummy_cam; + + int failures = 0; + + // --- Test 1: yolo image --- + { + std::cout << "\n=== Test 1: " << yolo_path << " ===" << std::endl; + engines::OpenVocab engine(config); + engine.enable(dummy_cam); + if (!run_test(engine, yolo_path)) ++failures; + } + + // --- Test 2: yoloworld image (optionally with a different model) --- + { + std::cout << "\n=== Test 2: " << yoloworld_path << " ===" << std::endl; + engines::OpenVocabConfig config2 = config; + const char* world_model = std::getenv("ZED_ONNX_WORLD_MODEL_PATH"); + if (world_model && *world_model) + config2.model_path = world_model; + + engines::OpenVocab engine(config2); + engine.enable(dummy_cam); + if (!run_test(engine, yoloworld_path)) ++failures; + } + + return failures > 0 ? 1 : 0; +} diff --git a/src/util/clip_tokenizer.h b/src/util/clip_tokenizer.h new file mode 100644 index 0000000..9834a2b --- /dev/null +++ b/src/util/clip_tokenizer.h @@ -0,0 +1,226 @@ +#pragma once +// Minimal CLIP ViT-B/32 BPE tokenizer — header-only, C++17, no external deps. +// +// Loads clip_vocab.txt (token\tid, tab-separated) and clip_merges.txt +// (a b, space-separated BPE pairs) produced by export_clip_text_encoder.py. +// Returns 77 int64 token IDs compatible with clip_text_encoder.onnx. + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace util { + +class ClipTokenizer { +public: + static constexpr int SEQ_LEN = 77; + static constexpr int64_t BOS = 49406; + static constexpr int64_t EOS = 49407; + static constexpr int64_t PAD = 0; + + void load(const std::string& vocab_path, const std::string& merges_path) { + load_vocab(vocab_path); + load_merges(merges_path); + init_byte_encoder(); + } + + bool is_loaded() const { return !vocab_.empty(); } + + // Tokenize a label string. Returns exactly SEQ_LEN int64 token IDs. + std::vector tokenize(const std::string& text) const { + std::vector ids; + ids.reserve(SEQ_LEN); + ids.push_back(BOS); + + std::string lower = text; + std::transform(lower.begin(), lower.end(), lower.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + + // Split on whitespace — sufficient for typical detection label names + std::istringstream ss(lower); + std::string word; + while (ss >> word) { + if (static_cast(ids.size()) >= SEQ_LEN - 1) break; + + // Byte-encode the word via the GPT-2 byte encoder + std::string encoded; + encoded.reserve(word.size() * 3); + for (unsigned char c : word) { + auto it = byte_encoder_.find(c); + if (it != byte_encoder_.end()) + encoded += utf8_of(it->second); + } + if (encoded.empty()) continue; + + // Apply BPE + auto tokens = bpe_encode(encoded); + for (const auto& tok : tokens) { + if (static_cast(ids.size()) >= SEQ_LEN - 1) break; + auto vit = vocab_.find(tok); + if (vit != vocab_.end()) + ids.push_back(static_cast(vit->second)); + } + } + + ids.push_back(EOS); + ids.resize(SEQ_LEN, PAD); + return ids; + } + +private: + std::unordered_map vocab_; // token string → id + std::unordered_map merge_ranks_; // "a b" → rank + std::unordered_map byte_encoder_; // byte → unicode codepoint + + // ----------------------------------------------------------------- + // Byte encoder — matches Python's bytes_to_unicode() in OpenAI CLIP + // ----------------------------------------------------------------- + void init_byte_encoder() { + // Printable ASCII + Latin-1 supplement ranges that map to themselves + std::vector bs; + for (int i = '!'; i <= '~'; i++) bs.push_back(i); + for (int i = 0xA1; i <= 0xAC; i++) bs.push_back(i); + for (int i = 0xAE; i <= 0xFF; i++) bs.push_back(i); + + std::vector cs; + cs.reserve(bs.size()); + for (int b : bs) cs.push_back(static_cast(b)); + + // Remaining bytes map to codepoints starting at U+0100 + int n = 0; + for (int b = 0; b < 256; b++) { + if (std::find(bs.begin(), bs.end(), b) == bs.end()) { + bs.push_back(b); + cs.push_back(256u + static_cast(n)); + ++n; + } + } + + byte_encoder_.clear(); + for (size_t i = 0; i < bs.size(); i++) + byte_encoder_[static_cast(bs[i])] = cs[i]; + } + + // Encode a unicode codepoint as a UTF-8 std::string + static std::string utf8_of(uint32_t cp) { + std::string s; + if (cp < 0x80u) { + s += static_cast(cp); + } else if (cp < 0x800u) { + s += static_cast(0xC0u | (cp >> 6)); + s += static_cast(0x80u | (cp & 0x3Fu)); + } else if (cp < 0x10000u) { + s += static_cast(0xE0u | (cp >> 12)); + s += static_cast(0x80u | ((cp >> 6) & 0x3Fu)); + s += static_cast(0x80u | (cp & 0x3Fu)); + } else { + s += static_cast(0xF0u | (cp >> 18)); + s += static_cast(0x80u | ((cp >> 12) & 0x3Fu)); + s += static_cast(0x80u | ((cp >> 6) & 0x3Fu)); + s += static_cast(0x80u | (cp & 0x3Fu)); + } + return s; + } + + // ----------------------------------------------------------------- + // BPE + // ----------------------------------------------------------------- + + // Split a UTF-8 string into individual Unicode character substrings + static std::vector split_utf8_chars(const std::string& s) { + std::vector chars; + for (size_t i = 0; i < s.size(); ) { + unsigned char c = static_cast(s[i]); + int len = (c < 0x80) ? 1 : (c < 0xE0) ? 2 : (c < 0xF0) ? 3 : 4; + chars.push_back(s.substr(i, len)); + i += len; + } + return chars; + } + + // BPE-encode a single byte-encoded word, returning the token strings + std::vector bpe_encode(const std::string& word_str) const { + std::vector word = split_utf8_chars(word_str); + if (word.size() <= 1) return word; + + while (true) { + // Find the pair with the lowest merge rank + int best_rank = INT_MAX; + int best_i = -1; + for (int i = 0; i + 1 < static_cast(word.size()); i++) { + auto it = merge_ranks_.find(word[i] + " " + word[i + 1]); + if (it != merge_ranks_.end() && it->second < best_rank) { + best_rank = it->second; + best_i = i; + } + } + if (best_i < 0) break; + + // Apply merge at best_i + std::string merged = word[best_i] + word[best_i + 1]; + word.erase(word.begin() + best_i, word.begin() + best_i + 2); + word.insert(word.begin() + best_i, std::move(merged)); + } + + return word; + } + + // ----------------------------------------------------------------- + // File loading + // ----------------------------------------------------------------- + + void load_vocab(const std::string& path) { + std::ifstream f(path); + if (!f.is_open()) + throw std::runtime_error("ClipTokenizer: cannot open vocab file: " + path); + + std::string line; + while (std::getline(f, line)) { + auto tab = line.rfind('\t'); + if (tab == std::string::npos) continue; + std::string token = unescape(line.substr(0, tab)); + int id = std::stoi(line.substr(tab + 1)); + vocab_[token] = id; + } + } + + void load_merges(const std::string& path) { + std::ifstream f(path); + if (!f.is_open()) + throw std::runtime_error("ClipTokenizer: cannot open merges file: " + path); + + std::string line; + int rank = 0; + while (std::getline(f, line)) { + if (line.empty() || line[0] == '#') continue; + auto sp = line.find(' '); + if (sp == std::string::npos) continue; + merge_ranks_[line] = rank++; // key is "a b" as written in the file + } + } + + static std::string unescape(const std::string& s) { + std::string r; + r.reserve(s.size()); + for (size_t i = 0; i < s.size(); i++) { + if (s[i] == '\\' && i + 1 < s.size()) { + switch (s[++i]) { + case '\\': r += '\\'; break; + case 'n': r += '\n'; break; + case 't': r += '\t'; break; + default: r += '\\'; r += s[i]; break; + } + } else { + r += s[i]; + } + } + return r; + } +}; + +} // namespace util diff --git a/test/assets/.gitkeep b/test/assets/.gitkeep new file mode 100644 index 0000000..e69de29