refactor open_vocab into yolo_onnx and yoloworld_onnx engines
Split the monolithic open_vocab engine into two focused engines: yolo_onnx (closed-vocabulary, any single-input YOLO ONNX) and yoloworld_onnx (open-vocabulary YOLO-World with CLIP encoder). Shared NMS/letterbox/fuse helpers extracted to yolo_helpers.h. Add export_yolo.py script and separate test binaries for each engine.
This commit is contained in:
cristhian aguilera
@@ -34,14 +34,18 @@ Camera and detection behavior is configured entirely through environment variabl
|
||||
- **`src/zed/camera.cc`** — Implementation. SDK enum parsers live in an anonymous namespace.
|
||||
- **`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/engines/yolo_helpers.h`** — Header-only shared helpers used by both ONNX engines: `RawDetection`, `nms()`, `letterbox()`, `fuse_bbox()`, `backproject()`, `split_labels()`. Lives in `engines::detail` namespace.
|
||||
- **`src/engines/yolo_onnx.h/cc`** — `engines::YoloOnnx` / `YoloOnnxConfig`: closed-vocabulary detection for any single-input YOLO ONNX model (YOLOv8, YOLOv11, YOLOv5, etc.). Auto-detects output layout (`NC_FIRST` / `NC_LAST` / `NC_LAST_OBJ`) from tensor shape; override with `ZED_ONNX_OUTPUT_FORMAT`. Labels loaded from `ZED_ONNX_LABELS` or `ZED_ONNX_LABELS_FILE`.
|
||||
- **`src/engines/yoloworld_onnx.h/cc`** — `engines::YoloWorldOnnx` / `YoloWorldOnnxConfig`: open-vocabulary detection via YOLO-World ONNX (two inputs: `images` + `text_feats`). Reads `ZED_ONNX_LABELS` on every tick; when labels change, re-tokenizes and re-encodes via the CLIP text encoder ONNX session, passing `(1, N, 512)` L2-normalized embeddings to the visual model. Output is always NC_FIRST `(1, 4+N, anchors)`.
|
||||
- **`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<engines::Engine>`. Parses `ZED_PC_MODE` (`FULL`/`HALF`/`ROI`) once before the loop. On `tick` input: calls `camera.grab()` and `detector->detect()`, then sends `image_bgr`, `point_cloud`, `camera_info`, and `objects` outputs. Point cloud output varies by mode: `FULL` sends the raw `sl::Mat` pointer (no copy); `HALF` downsamples by 2 in each dimension; `ROI` emits only valid points within detected bounding boxes as a flat `[N, 4]` array. On `update_vocabulary` input: decodes payload as a UTF-8 label string and calls `detector->set_labels()` for immediate re-encoding. Handles reconnection on grab failure.
|
||||
- **`src/main.cc`** — Dora event loop. Selects engine via `ZED_ENGINE` env var (`yolo_onnx`/`onnx`, `yoloworld_onnx`/`open_vocab`, or `native`), holds a `std::unique_ptr<engines::Engine>`. Parses `ZED_PC_MODE` (`FULL`/`HALF`/`ROI`) once before the loop. On `tick` input: calls `camera.grab()` and `detector->detect()`, then sends `image_bgr`, `point_cloud`, `camera_info`, and `objects` outputs. Point cloud output varies by mode: `FULL` sends the raw `sl::Mat` pointer (no copy); `HALF` downsamples by 2 in each dimension; `ROI` emits only valid points within detected bounding boxes as a flat `[N, 4]` array. On `update_vocabulary` input: decodes payload as a UTF-8 label string and calls `detector->set_labels()` for immediate re-encoding. 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_yolo.py`** — PEP 723 uv script. Exports a standard YOLOv11 model to ONNX for closed-vocabulary detection. Run: `uv run scripts/export_yolo.py --model n --output-dir models/`.
|
||||
- **`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.
|
||||
- **`src/test_yolo_onnx.cc`** — Standalone C++ test binary (`bin/test_yolo_onnx`). Exercises `YoloOnnx` with an image (default `test/assets/yolo.png`). No dora and no live camera required. Uses the same env vars as the main binary. Saves an annotated result image alongside the input.
|
||||
- **`src/test_yoloworld_onnx.cc`** — Standalone C++ test binary (`bin/test_yoloworld_onnx`). Exercises `YoloWorldOnnx` with an image (default `test/assets/yoloworld.png`). Same interface as `test_yolo_onnx`.
|
||||
|
||||
## Key APIs
|
||||
|
||||
|
||||
+28
-8
@@ -24,7 +24,8 @@ add_executable(dora_zed
|
||||
src/main.cc
|
||||
src/zed/camera.cc
|
||||
src/engines/zed_native.cc
|
||||
src/engines/open_vocab.cc
|
||||
src/engines/yolo_onnx.cc
|
||||
src/engines/yoloworld_onnx.cc
|
||||
${node_bridge}
|
||||
)
|
||||
add_dependencies(dora_zed Dora_cxx)
|
||||
@@ -45,21 +46,40 @@ target_link_libraries(dora_zed dora_node_api_cxx ${ZED_LIBRARIES} ${OpenCV_LIBS}
|
||||
|
||||
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
|
||||
# Standalone detector tests — no dora, no live camera required
|
||||
|
||||
add_executable(test_yolo_onnx
|
||||
src/test_yolo_onnx.cc
|
||||
src/engines/yolo_onnx.cc
|
||||
)
|
||||
target_include_directories(test_open_vocab PRIVATE
|
||||
target_include_directories(test_yolo_onnx 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
|
||||
target_link_libraries(test_yolo_onnx PRIVATE
|
||||
${ZED_LIBRARIES}
|
||||
${OpenCV_LIBS}
|
||||
${ONNXRUNTIME_LIBRARIES}
|
||||
)
|
||||
install(TARGETS test_open_vocab DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/bin)
|
||||
install(TARGETS test_yolo_onnx DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/bin)
|
||||
|
||||
add_executable(test_yoloworld_onnx
|
||||
src/test_yoloworld_onnx.cc
|
||||
src/engines/yoloworld_onnx.cc
|
||||
)
|
||||
target_include_directories(test_yoloworld_onnx PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/src
|
||||
${ZED_INCLUDE_DIRS}
|
||||
${OpenCV_INCLUDE_DIRS}
|
||||
${CUDAToolkit_INCLUDE_DIRS}
|
||||
${ONNXRUNTIME_INCLUDE_DIRS}
|
||||
)
|
||||
target_link_libraries(test_yoloworld_onnx PRIVATE
|
||||
${ZED_LIBRARIES}
|
||||
${OpenCV_LIBS}
|
||||
${ONNXRUNTIME_LIBRARIES}
|
||||
)
|
||||
install(TARGETS test_yoloworld_onnx DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/bin)
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
# dora-zed
|
||||
|
||||
Dora C++ node for ZED stereo cameras. Captures frames and publishes BGR images, XYZRGBA point clouds, camera intrinsics, and 3D object detections.
|
||||
Dora node for ZED stereo cameras. Publishes BGR images, point clouds, camera intrinsics, and 3D object detections.
|
||||
|
||||
## Usage via `git:` in a dataflow
|
||||
## Usage
|
||||
|
||||
Add to your dataflow:
|
||||
|
||||
```yaml
|
||||
nodes:
|
||||
- id: zed-camera
|
||||
git: ssh://git@gitea.liria.cl:5001/liria/dora-zed.git
|
||||
git_checkout: v0.1.0
|
||||
build: mkdir -p build && cd build && cmake .. && make install
|
||||
path: bin/dora_zed
|
||||
inputs:
|
||||
@@ -19,180 +22,192 @@ nodes:
|
||||
- objects
|
||||
```
|
||||
|
||||
Then run:
|
||||
See [`dataflow.yml`](dataflow.yml) for a complete runnable example with a camera display node.
|
||||
|
||||
```bash
|
||||
dora build dataflow.yml
|
||||
dora up
|
||||
dora start dataflow.yml
|
||||
```
|
||||
## Dependencies
|
||||
|
||||
| Dependency | Version | Notes |
|
||||
|---|---|---|
|
||||
| [CMake](https://cmake.org/download/) | 3.21+ | |
|
||||
| [ZED SDK](https://www.stereolabs.com/developers/release) | 5.x | Bundles OpenCV; default path `/usr/local/zed` |
|
||||
| [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) | — | Required by ZED SDK |
|
||||
| [Rust toolchain](https://rustup.rs) | stable | Builds dora-rs C++ bindings |
|
||||
| [ONNX Runtime](https://github.com/microsoft/onnxruntime/releases) | 1.20+ | Required only for `open_vocab` engine |
|
||||
|
||||
### Installation scripts
|
||||
|
||||
| Platform | Script |
|
||||
|---|---|
|
||||
| Ubuntu 24.04 x86_64 | `sudo scripts/install_linux_x64.sh` |
|
||||
| NVIDIA Jetson (ARM64) | `sudo scripts/install_linux_arm.sh` |
|
||||
| Windows x64 (PowerShell admin) | `scripts/install_windows.ps1` |
|
||||
|
||||
> **Jetson:** No official ONNX Runtime GPU binary for ARM64. The `native` engine works without it. The `open_vocab` engine requires [building from source](https://onnxruntime.ai/docs/build/eps.html#tensorrt).
|
||||
>
|
||||
> **Windows:** Install ZED SDK manually from [stereolabs.com](https://www.stereolabs.com/developers/release), then pass the ONNX Runtime path: `cmake .. -DCMAKE_PREFIX_PATH=C:\onnxruntime`
|
||||
|
||||
## Local build
|
||||
|
||||
```bash
|
||||
git clone ssh://git@gitea.liria.cl:5001/liria/dora-zed.git
|
||||
cd dora-zed
|
||||
mkdir -p build && cd build
|
||||
cmake ..
|
||||
make install
|
||||
mkdir -p build && cd build && cmake .. && make install
|
||||
```
|
||||
|
||||
The binary is installed to `bin/dora_zed` relative to the repository root.
|
||||
Binary installed to `bin/dora_zed`.
|
||||
|
||||
## Dependencies
|
||||
## Configuration
|
||||
|
||||
| Dependency | Minimum version | Notes |
|
||||
|---|---|---|
|
||||
| [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 |
|
||||
All behavior is controlled via environment variables.
|
||||
|
||||
## 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
|
||||
### Camera
|
||||
|
||||
| Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `ZED_RESOLUTION` | `720` | Camera resolution: `720`, `1080`, or `2000` |
|
||||
| `ZED_FPS` | `15` | Camera frame rate |
|
||||
| `ZED_DEPTH_MODE` | `NEURAL` | Depth mode: `PERFORMANCE`, `QUALITY`, `NEURAL`, or `NONE` |
|
||||
| `ZED_DEPTH_MIN_MM` | `10` | Minimum depth distance in millimeters |
|
||||
| `ZED_DEPTH_MAX_MM` | `500` | Maximum depth distance in millimeters |
|
||||
| `ZED_RESOLUTION` | `720` | `720`, `1080`, or `2000` |
|
||||
| `ZED_FPS` | `15` | Frame rate |
|
||||
| `ZED_DEPTH_MODE` | `NEURAL` | `PERFORMANCE`, `QUALITY`, `NEURAL`, or `NONE` |
|
||||
| `ZED_DEPTH_MIN_MM` | `10` | Minimum depth (mm) |
|
||||
| `ZED_DEPTH_MAX_MM` | `500` | Maximum depth (mm) |
|
||||
| `ZED_DEPTH_FILL` | `false` | Enable depth fill mode |
|
||||
| `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_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 |
|
||||
| `ZED_PC_MODE` | `FULL` | Point cloud output mode: `FULL` (full HxW grid), `HALF` (stride-2 downsample, ~4× smaller), or `ROI` (valid points within detected bounding boxes only, flat list) |
|
||||
| `ZED_FLIP` | `ON` | `ON` or `OFF` |
|
||||
| `ZED_WARMUP_FRAMES` | `30` | Frames to discard on startup |
|
||||
| `ZED_RETRY_INTERVAL_SECONDS` | `5` | Reconnection interval (s) |
|
||||
|
||||
### Detection engine
|
||||
|
||||
Set `ZED_ENGINE` to select the detection backend:
|
||||
|
||||
| `ZED_ENGINE` | Custom model? | Classes |
|
||||
|---|---|---|
|
||||
| `native` (default) | No — ZED SDK built-in TensorRT | Fixed: 7 categories |
|
||||
| `yolo_onnx` (alias: `onnx`) | Yes — any single-input YOLO ONNX (v5/v8/v11) | Fixed labels from file or env var |
|
||||
| `yoloworld_onnx` (alias: `open_vocab`) | Yes — YOLO-World two-input ONNX | Defined at runtime via CLIP |
|
||||
|
||||
#### `native` engine
|
||||
|
||||
Uses ZED SDK's built-in TensorRT models — no model files to download or configure. `ZED_DETECTION_MODEL` selects the speed/accuracy variant, not a file path.
|
||||
|
||||
Detects 7 fixed categories: PERSON, VEHICLE, BAG, ANIMAL, ELECTRONICS, FRUIT_VEGETABLE, SPORT.
|
||||
|
||||
| Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `ZED_DETECTION_MODEL` | `MULTI_CLASS_BOX_MEDIUM` | `MULTI_CLASS_BOX_FAST`, `MULTI_CLASS_BOX_MEDIUM`, or `MULTI_CLASS_BOX_ACCURATE` |
|
||||
| `ZED_DETECTION_CONFIDENCE` | `40` | Confidence threshold (0–100) |
|
||||
| `ZED_DETECTION_TRACKING` | `false` | Enable object tracking |
|
||||
|
||||
#### `yolo_onnx` engine (closed-vocabulary)
|
||||
|
||||
For any single-input YOLO ONNX model with fixed class names (YOLOv8, YOLOv11, YOLOv5, custom-trained).
|
||||
|
||||
Export a YOLOv11 model once:
|
||||
|
||||
```bash
|
||||
uv run scripts/export_yolo.py --model n --output-dir models/
|
||||
```
|
||||
|
||||
Then run:
|
||||
|
||||
```bash
|
||||
ZED_ENGINE=yolo_onnx \
|
||||
ZED_ONNX_MODEL_PATH=models/yolo11n_i640.onnx \
|
||||
ZED_ONNX_LABELS_FILE=models/coco80.txt \
|
||||
dora start dataflow.yml
|
||||
```
|
||||
|
||||
**Required:**
|
||||
|
||||
| Variable | Description |
|
||||
|---|---|
|
||||
| `ZED_ONNX_MODEL_PATH` | Path to YOLO ONNX model |
|
||||
| `ZED_ONNX_LABELS` or `ZED_ONNX_LABELS_FILE` | Class names (inline CSV or one-per-line file) |
|
||||
|
||||
**Optional:**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `ZED_ONNX_CONFIDENCE` | `0.25` | Confidence threshold (0.0–1.0) |
|
||||
| `ZED_ONNX_NMS_THRESHOLD` | `0.45` | NMS IoU threshold |
|
||||
| `ZED_ONNX_INPUT_SIZE` | `640` | Model input size (px) |
|
||||
| `ZED_ONNX_OUTPUT_FORMAT` | `auto` | Output layout: `auto`, `nc_first`/`yolov8`/`yolov11`, `nc_last`, `nc_last_obj`/`yolov5`. `auto` detects from tensor shape. |
|
||||
|
||||
#### `yoloworld_onnx` engine (open-vocabulary)
|
||||
|
||||
For YOLO-World two-input ONNX models with a CLIP text encoder. Labels are defined at runtime — no recompilation required. Use `ZED_ENGINE=open_vocab` as a legacy alias.
|
||||
|
||||
Export models once:
|
||||
|
||||
```bash
|
||||
uv run scripts/export_yoloworld.py --model s --num-classes 10 --output-dir models/
|
||||
uv run scripts/export_clip_text_encoder.py --output-dir models/
|
||||
```
|
||||
|
||||
Then run:
|
||||
|
||||
```bash
|
||||
ZED_ENGINE=yoloworld_onnx \
|
||||
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,chair" \
|
||||
dora start dataflow.yml
|
||||
```
|
||||
|
||||
**Required:**
|
||||
|
||||
| Variable | Description |
|
||||
|---|---|
|
||||
| `ZED_ONNX_MODEL_PATH` | Path to YOLO-World ONNX model |
|
||||
| `ZED_ONNX_TEXT_ENCODER_PATH` | Path to `clip_text_encoder.onnx` |
|
||||
| `ZED_ONNX_CLIP_VOCAB` | Path to `clip_vocab.txt` |
|
||||
| `ZED_ONNX_CLIP_MERGES` | Path to `clip_merges.txt` |
|
||||
| `ZED_ONNX_LABELS` | Comma-separated labels, e.g. `person,car,chair` |
|
||||
|
||||
**Optional:**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `ZED_ONNX_CONFIDENCE` | `0.25` | Confidence threshold (0.0–1.0) |
|
||||
| `ZED_ONNX_NMS_THRESHOLD` | `0.45` | NMS IoU threshold |
|
||||
| `ZED_ONNX_INPUT_SIZE` | `640` | Model input size (px) |
|
||||
|
||||
Labels can be updated at runtime without restarting — send a new comma-separated string to the `update_vocabulary` input.
|
||||
|
||||
### Point cloud
|
||||
|
||||
| `ZED_PC_MODE` | Shape | Description |
|
||||
|---|---|---|
|
||||
| `FULL` (default) | `[H, W, 4]` | Full resolution grid (~14.7 MB/frame at 720p) |
|
||||
| `HALF` | `[H/2, W/2, 4]` | Stride-2 downsample (~4× smaller) |
|
||||
| `ROI` | `[N, 4]` | Valid points within detected bounding boxes only |
|
||||
|
||||
## Inputs
|
||||
|
||||
| Input | Description |
|
||||
|---|---|
|
||||
| `tick` | Required. Triggers a camera grab and detection on each message. Connect to `dora/timer/millis/N`. |
|
||||
| `update_vocabulary` | Optional (`open_vocab` engine only). Payload: raw UTF-8 bytes of a comma-separated label string (e.g. `"person,cup,chair"`). Immediately re-encodes CLIP embeddings; the next `tick` uses the new labels. |
|
||||
| `tick` | Triggers a camera grab and detection. Connect to `dora/timer/millis/N`. |
|
||||
| `update_vocabulary` | (`open_vocab` only) UTF-8 comma-separated label string. Re-encodes CLIP embeddings immediately. |
|
||||
|
||||
## Outputs
|
||||
|
||||
### `image_bgr`
|
||||
| Output | Type | Description |
|
||||
|---|---|---|
|
||||
| `image_bgr` | `uint8 [H, W, 3]` | BGR image. Metadata: `shape`, `encoding=bgr8`, intrinsics, `timestamp_ns`. |
|
||||
| `point_cloud` | `float32 [*, 4]` | XYZRGBA point cloud. Shape depends on `ZED_PC_MODE`. Metadata: `units=mm`, intrinsics. |
|
||||
| `camera_info` | `float32 [4]` | `[fx, fy, cx, cy]`. Metadata: `distortion` coefficients. |
|
||||
| `objects` | `float32 [N, 37]` | Detected objects. See below. |
|
||||
|
||||
BGR image as raw bytes. Metadata includes `shape` (HxWx3), `encoding` (`bgr8`), camera intrinsics (`fx`, `fy`, `cx`, `cy`), `distortion`, `frame_id`, and `timestamp_ns`.
|
||||
### `objects` layout
|
||||
|
||||
### `point_cloud`
|
||||
|
||||
XYZRGBA point cloud as float32 bytes. Controlled by `ZED_PC_MODE`:
|
||||
|
||||
| Mode | Shape | Layout | Description |
|
||||
|---|---|---|---|
|
||||
| `FULL` | `[H, W, 4]` | `HWC` | Full resolution grid (default). ~14.7 MB/frame at 720p |
|
||||
| `HALF` | `[H/2, W/2, 4]` | `HWC` | Stride-2 downsample (~4× smaller) |
|
||||
| `ROI` | `[N, 4]` | `NC` | Valid points within detected object bounding boxes only. `N` varies by scene. Metadata includes `mode=roi` |
|
||||
|
||||
Metadata includes `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:
|
||||
Each of the N rows contains 37 floats:
|
||||
|
||||
| Index | Field | Description |
|
||||
|---|---|---|
|
||||
| 0 | `label_id` | Object class (0=PERSON, 1=VEHICLE, 2=BAG, 3=ANIMAL, 4=ELECTRONICS, 5=FRUIT_VEGETABLE, 6=SPORT) |
|
||||
| 1 | `confidence` | Detection confidence (0-100) |
|
||||
| 2-4 | `x, y, z` | 3D center position in mm |
|
||||
| 5-12 | `bb2d_*` | 2D bounding box: 4 corners × 2 pixel coords |
|
||||
| 13-36 | `bb3d_*` | 3D bounding box: 8 corners × 3 world coords (mm) |
|
||||
| 0 | `label_id` | Class index |
|
||||
| 1 | `confidence` | Score (0–100) |
|
||||
| 2–4 | `x, y, z` | 3D center (mm) |
|
||||
| 5–12 | `bb2d` | 2D bounding box: 4 corners × (x, y) px |
|
||||
| 13–36 | `bb3d` | 3D bounding box: 8 corners × (x, y, z) mm |
|
||||
|
||||
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.
|
||||
Metadata includes `engine`, `label_names`, `image_width`, `image_height`, and `timestamp_ns`. An empty array `[0, 37]` is sent when no objects are detected.
|
||||
|
||||
+8
-3
@@ -1,13 +1,18 @@
|
||||
nodes:
|
||||
- id: zed-camera
|
||||
git: ssh://git@gitea.liria.cl:5001/liria/dora-zed.git
|
||||
git_checkout: v0.1.0
|
||||
build: mkdir -p build && cd build && cmake .. && make install
|
||||
path: bin/dora_zed
|
||||
inputs:
|
||||
tick: dora/timer/millis/100
|
||||
# update_vocabulary: vocabulary-node/labels # optional: send new comma-separated labels at runtime
|
||||
outputs:
|
||||
- image_bgr
|
||||
- point_cloud
|
||||
- camera_info
|
||||
- objects
|
||||
|
||||
- id: plot
|
||||
git: https://github.com/dora-rs/dora-hub.git
|
||||
build: pip install ./node-hub/opencv-plot
|
||||
path: opencv-plot
|
||||
inputs:
|
||||
image: zed-camera/image_bgr
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["ultralytics>=8.3"]
|
||||
# ///
|
||||
"""Export a standard YOLOv11 model to ONNX for closed-vocabulary detection.
|
||||
|
||||
The exported model has one input:
|
||||
images (1, 3, imgsz, imgsz) float32
|
||||
|
||||
Class names are fixed at training time. Provide them at runtime via
|
||||
ZED_ONNX_LABELS or ZED_ONNX_LABELS_FILE.
|
||||
|
||||
Usage:
|
||||
uv run scripts/export_yolo.py
|
||||
uv run scripts/export_yolo.py --model s --output-dir models/
|
||||
uv run scripts/export_yolo.py --model m --imgsz 1280
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
MODELS = {
|
||||
"n": "yolo11n",
|
||||
"s": "yolo11s",
|
||||
"m": "yolo11m",
|
||||
"l": "yolo11l",
|
||||
"x": "yolo11x",
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(
|
||||
description="Export YOLOv11 to ONNX for closed-vocabulary detection"
|
||||
)
|
||||
p.add_argument(
|
||||
"--model", "-m",
|
||||
choices=list(MODELS),
|
||||
default="n",
|
||||
help="Model size: n (~6 MB), s (~18 MB), m (~38 MB), l (~49 MB), x (~109 MB). Default: n",
|
||||
)
|
||||
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_output(onnx_path: Path):
|
||||
try:
|
||||
import onnx
|
||||
model = onnx.load(str(onnx_path))
|
||||
inputs = [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim])
|
||||
for i in model.graph.input]
|
||||
outputs = [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim])
|
||||
for o in model.graph.output]
|
||||
print(f" Inputs: {inputs}")
|
||||
print(f" Outputs: {outputs}")
|
||||
except ImportError:
|
||||
print(" (install 'onnx' to verify model structure)")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
except ImportError:
|
||||
print("Error: ultralytics not installed. Run: uv run scripts/export_yolo.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}_i{args.imgsz}.onnx"
|
||||
dst = out_dir / out_name
|
||||
|
||||
print(f"Model: {model_name}.pt")
|
||||
print(f"Input size: {args.imgsz}")
|
||||
print(f"Output: {dst}")
|
||||
print()
|
||||
|
||||
model = YOLO(f"{model_name}.pt")
|
||||
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_output(dst)
|
||||
|
||||
print(f"\nExported: {dst.resolve()}")
|
||||
print(f"\nRun with (COCO labels example):")
|
||||
print(f" ZED_ENGINE=onnx \\")
|
||||
print(f" ZED_ONNX_MODEL_PATH={dst.resolve()} \\")
|
||||
print(f" ZED_ONNX_LABELS_FILE=models/coco80.txt \\")
|
||||
print(f" dora start dataflow.yml")
|
||||
print(f"\nOr with inline labels:")
|
||||
print(f" ZED_ONNX_LABELS=\"person,car,chair\" \\")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,592 +0,0 @@
|
||||
#include "engines/open_vocab.h"
|
||||
#include "util/clip_tokenizer.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <onnxruntime_cxx_api.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <numeric>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
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<RawDetection> nms(std::vector<RawDetection> dets, float threshold) {
|
||||
std::sort(dets.begin(), dets.end(),
|
||||
[](const RawDetection& a, const RawDetection& b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
std::vector<bool> suppressed(dets.size(), false);
|
||||
std::vector<RawDetection> 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<float>(size) / src.cols;
|
||||
float scale_h = static_cast<float>(size) / src.rows;
|
||||
scale = std::min(scale_w, scale_h);
|
||||
int new_w = static_cast<int>(std::round(src.cols * scale));
|
||||
int new_h = static_cast<int>(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<int>(pad_x), static_cast<int>(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<int>(pc.getWidth());
|
||||
int h = static_cast<int>(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<float> 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<std::string> split_labels(const std::string& s) {
|
||||
std::vector<std::string> 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<Ort::Session> 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<Ort::Session> 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<float> 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<RawDetection> run_inference(const OpenVocabConfig& config,
|
||||
const cv::Mat& image_bgr);
|
||||
};
|
||||
|
||||
// --- Impl::run_inference ---
|
||||
|
||||
std::vector<RawDetection> 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<float> 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<cv::Vec3f>(h, w)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
std::vector<int64_t> image_shape{1, 3, sz, sz};
|
||||
Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
|
||||
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<Ort::Value> outputs;
|
||||
|
||||
try {
|
||||
if (has_text_input) {
|
||||
std::vector<int64_t> text_shape{1, num_classes, 512};
|
||||
Ort::Value text_tensor = Ort::Value::CreateTensor<float>(
|
||||
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<int>(out_shape[2]);
|
||||
int nc = num_classes;
|
||||
if (nc <= 0) return {};
|
||||
|
||||
const float* out = outputs[0].GetTensorData<float>();
|
||||
|
||||
std::vector<RawDetection> 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<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(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<Impl>()) {}
|
||||
|
||||
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<Ort::Session>(
|
||||
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<int>(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<Ort::Session>(
|
||||
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<size_t>(N) * 512, 0.0f);
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
int encode_n = std::min(static_cast<int>(label_vec.size()), N);
|
||||
for (int i = 0; i < encode_n; i++) {
|
||||
auto tokens = impl_->tokenizer.tokenize(label_vec[i]);
|
||||
|
||||
std::vector<int64_t> input_shape{1, util::ClipTokenizer::SEQ_LEN};
|
||||
Ort::Value token_tensor = Ort::Value::CreateTensor<int64_t>(
|
||||
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<Ort::Value> 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<float>();
|
||||
|
||||
// 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<int>(img_mat.getWidth());
|
||||
int orig_h = static_cast<int>(img_mat.getHeight());
|
||||
|
||||
// RGBA → BGR for run_inference
|
||||
cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
|
||||
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<int>(det.x1), static_cast<int>(det.y1),
|
||||
static_cast<int>(det.x2), static_cast<int>(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);
|
||||
}
|
||||
}
|
||||
|
||||
void OpenVocab::set_labels(const std::string& labels) {
|
||||
config_.labels = labels;
|
||||
encode_labels(labels);
|
||||
}
|
||||
|
||||
const std::vector<DetectedObject>& OpenVocab::objects() const {
|
||||
return detected_;
|
||||
}
|
||||
|
||||
std::string OpenVocab::label_names() const {
|
||||
return impl_->cached_labels.empty() ? config_.labels : impl_->cached_labels;
|
||||
}
|
||||
|
||||
} // namespace engines
|
||||
@@ -0,0 +1,153 @@
|
||||
#pragma once
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace engines::detail {
|
||||
|
||||
// --- NMS ---
|
||||
|
||||
struct RawDetection {
|
||||
float x1, y1, x2, y2;
|
||||
float score;
|
||||
int class_id;
|
||||
};
|
||||
|
||||
inline 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;
|
||||
}
|
||||
|
||||
inline std::vector<RawDetection> nms(std::vector<RawDetection> dets, float threshold) {
|
||||
std::sort(dets.begin(), dets.end(),
|
||||
[](const RawDetection& a, const RawDetection& b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
std::vector<bool> suppressed(dets.size(), false);
|
||||
std::vector<RawDetection> 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 ---
|
||||
|
||||
inline cv::Mat letterbox(const cv::Mat& src, int size,
|
||||
float& scale, float& pad_x, float& pad_y) {
|
||||
float scale_w = static_cast<float>(size) / src.cols;
|
||||
float scale_h = static_cast<float>(size) / src.rows;
|
||||
scale = std::min(scale_w, scale_h);
|
||||
int new_w = static_cast<int>(std::round(src.cols * scale));
|
||||
int new_h = static_cast<int>(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<int>(pad_x), static_cast<int>(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;
|
||||
};
|
||||
|
||||
inline FusionResult fuse_bbox(const sl::Mat& pc, int x1, int y1, int x2, int y2) {
|
||||
FusionResult result{};
|
||||
int w = static_cast<int>(pc.getWidth());
|
||||
int h = static_cast<int>(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<float> 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;
|
||||
}
|
||||
|
||||
inline 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;
|
||||
}
|
||||
|
||||
// --- Label parsing ---
|
||||
|
||||
inline std::vector<std::string> split_labels(const std::string& s) {
|
||||
std::vector<std::string> 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 engines::detail
|
||||
@@ -0,0 +1,373 @@
|
||||
#include "engines/yolo_onnx.h"
|
||||
#include "engines/yolo_helpers.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <onnxruntime_cxx_api.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace engines {
|
||||
|
||||
using namespace detail;
|
||||
|
||||
// --- Impl ---
|
||||
|
||||
struct YoloOnnx::Impl {
|
||||
Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "yolo_onnx"};
|
||||
Ort::SessionOptions session_options;
|
||||
std::unique_ptr<Ort::Session> session;
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
|
||||
std::string input_name;
|
||||
std::string output_name;
|
||||
int num_classes = 0;
|
||||
YoloOutputLayout resolved_layout = YoloOutputLayout::NC_FIRST;
|
||||
|
||||
std::vector<RawDetection> run_inference(const YoloOnnxConfig& config,
|
||||
const cv::Mat& image_bgr);
|
||||
};
|
||||
|
||||
// --- Impl::run_inference ---
|
||||
|
||||
std::vector<RawDetection> YoloOnnx::Impl::run_inference(
|
||||
const YoloOnnxConfig& config, const cv::Mat& image_bgr)
|
||||
{
|
||||
if (!session) return {};
|
||||
|
||||
int orig_w = image_bgr.cols;
|
||||
int orig_h = image_bgr.rows;
|
||||
|
||||
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<float> 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<cv::Vec3f>(h, w)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
std::vector<int64_t> image_shape{1, 3, sz, sz};
|
||||
Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info, input_data.data(), input_data.size(),
|
||||
image_shape.data(), image_shape.size());
|
||||
|
||||
const char* input_name_cstr = input_name.c_str();
|
||||
const char* output_name_cstr = output_name.c_str();
|
||||
|
||||
std::vector<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = session->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
&input_name_cstr, &image_tensor, 1,
|
||||
&output_name_cstr, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloOnnx: inference failed: " << e.what() << std::endl;
|
||||
return {};
|
||||
}
|
||||
|
||||
auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
|
||||
if (out_shape.size() < 3) return {};
|
||||
|
||||
const float* out = outputs[0].GetTensorData<float>();
|
||||
int nc = num_classes;
|
||||
if (nc <= 0) return {};
|
||||
|
||||
std::vector<RawDetection> candidates;
|
||||
|
||||
if (resolved_layout == YoloOutputLayout::NC_FIRST) {
|
||||
// (1, 4+N, anchors) — YOLOv8, YOLOv11
|
||||
int num_anchors = static_cast<int>(out_shape[2]);
|
||||
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<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
} else {
|
||||
// (1, anchors, 4+N) or (1, anchors, 5+N)
|
||||
int num_anchors = static_cast<int>(out_shape[1]);
|
||||
int row_stride = static_cast<int>(out_shape[2]);
|
||||
int cls_offset = (resolved_layout == YoloOutputLayout::NC_LAST_OBJ) ? 5 : 4;
|
||||
|
||||
for (int a = 0; a < num_anchors; a++) {
|
||||
const float* row = out + a * row_stride;
|
||||
float cx_lb = row[0], cy_lb = row[1], w_lb = row[2], h_lb = row[3];
|
||||
float obj_conf = (resolved_layout == YoloOutputLayout::NC_LAST_OBJ) ? row[4] : 1.0f;
|
||||
|
||||
float best_score = 0.0f;
|
||||
int best_class = 0;
|
||||
for (int c = 0; c < nc; c++) {
|
||||
float s = obj_conf * row[cls_offset + c];
|
||||
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<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
}
|
||||
|
||||
return nms(std::move(candidates), config.nms_threshold);
|
||||
}
|
||||
|
||||
// --- Config ---
|
||||
|
||||
YoloOnnxConfig YoloOnnxConfig::from_env() {
|
||||
YoloOnnxConfig 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.labels_file = util::get_env_str("ZED_ONNX_LABELS_FILE", c.labels_file);
|
||||
|
||||
std::string fmt = util::get_env_str("ZED_ONNX_OUTPUT_FORMAT", "");
|
||||
if (fmt == "nc_first" || fmt == "yolov8" || fmt == "yolov11")
|
||||
c.output_layout = YoloOutputLayout::NC_FIRST;
|
||||
else if (fmt == "nc_last")
|
||||
c.output_layout = YoloOutputLayout::NC_LAST;
|
||||
else if (fmt == "nc_last_obj" || fmt == "yolov5")
|
||||
c.output_layout = YoloOutputLayout::NC_LAST_OBJ;
|
||||
|
||||
return c;
|
||||
}
|
||||
|
||||
// --- YoloOnnx ---
|
||||
|
||||
YoloOnnx::YoloOnnx(const YoloOnnxConfig& config)
|
||||
: config_(config), impl_(std::make_unique<Impl>()) {}
|
||||
|
||||
YoloOnnx::~YoloOnnx() = default;
|
||||
|
||||
void YoloOnnx::enable(sl::Camera& /* camera */) {
|
||||
if (config_.model_path.empty()) {
|
||||
std::cerr << "YoloOnnx: ZED_ONNX_MODEL_PATH is not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
OrtCUDAProviderOptions cuda_opts{};
|
||||
cuda_opts.device_id = 0;
|
||||
try {
|
||||
impl_->session_options.AppendExecutionProvider_CUDA(cuda_opts);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloOnnx: CUDA unavailable (" << e.what() << "), using CPU." << std::endl;
|
||||
}
|
||||
impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "YoloOnnx: loading model (" << config_.model_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.model_path.c_str(), impl_->session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloOnnx: failed to load 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();
|
||||
if (num_inputs >= 2) {
|
||||
std::cerr << "YoloOnnx: model has " << num_inputs << " inputs — "
|
||||
"this looks like a YOLO-World model. Use ZED_ENGINE=yoloworld_onnx instead."
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
auto out_info = impl_->session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo();
|
||||
auto shape = out_info.GetShape();
|
||||
if (shape.size() == 3) {
|
||||
int64_t dim1 = shape[1];
|
||||
int64_t dim2 = shape[2];
|
||||
|
||||
YoloOutputLayout resolved;
|
||||
if (config_.output_layout == YoloOutputLayout::AUTO) {
|
||||
// dim1 < dim2 → NC_FIRST (e.g. 84 vs 8400 for YOLOv8/11)
|
||||
// dim1 > dim2 → NC_LAST_OBJ (e.g. 25200 vs 85 for YOLOv5)
|
||||
resolved = (dim1 < dim2) ? YoloOutputLayout::NC_FIRST
|
||||
: YoloOutputLayout::NC_LAST_OBJ;
|
||||
} else {
|
||||
resolved = config_.output_layout;
|
||||
}
|
||||
impl_->resolved_layout = resolved;
|
||||
|
||||
if (resolved == YoloOutputLayout::NC_FIRST) {
|
||||
impl_->num_classes = static_cast<int>(dim1) - 4;
|
||||
} else if (resolved == YoloOutputLayout::NC_LAST) {
|
||||
impl_->num_classes = static_cast<int>(dim2) - 4;
|
||||
} else { // NC_LAST_OBJ
|
||||
impl_->num_classes = static_cast<int>(dim2) - 5;
|
||||
}
|
||||
}
|
||||
|
||||
if (!config_.labels_file.empty() && config_.labels.empty()) {
|
||||
std::ifstream lf(config_.labels_file);
|
||||
if (!lf) {
|
||||
std::cerr << "YoloOnnx: cannot open ZED_ONNX_LABELS_FILE: "
|
||||
<< config_.labels_file << std::endl;
|
||||
} else {
|
||||
std::string line, joined;
|
||||
while (std::getline(lf, line)) {
|
||||
size_t l = line.find_first_not_of(" \t\r");
|
||||
size_t r = line.find_last_not_of(" \t\r");
|
||||
if (l == std::string::npos) continue;
|
||||
if (!joined.empty()) joined += ',';
|
||||
joined += line.substr(l, r - l + 1);
|
||||
}
|
||||
if (!joined.empty()) {
|
||||
config_.labels = joined;
|
||||
std::cout << "YoloOnnx: loaded labels from file ("
|
||||
<< config_.labels_file << ")" << std::endl;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::cout << "YoloOnnx: model loaded."
|
||||
<< " Classes=" << impl_->num_classes
|
||||
<< " Layout=" << (impl_->resolved_layout == YoloOutputLayout::NC_FIRST ? "NC_FIRST"
|
||||
: impl_->resolved_layout == YoloOutputLayout::NC_LAST ? "NC_LAST"
|
||||
: "NC_LAST_OBJ")
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
void YoloOnnx::detect(sl::Camera& camera) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
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<int>(img_mat.getWidth());
|
||||
int orig_h = static_cast<int>(img_mat.getHeight());
|
||||
|
||||
cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
|
||||
cv::Mat bgr;
|
||||
cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR);
|
||||
|
||||
auto kept = impl_->run_inference(config_, bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
FusionResult fusion = fuse_bbox(pc_mat,
|
||||
static_cast<int>(det.x1), static_cast<int>(det.y1),
|
||||
static_cast<int>(det.x2), static_cast<int>(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);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloOnnx::detect(const cv::Mat& image_bgr) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) 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;
|
||||
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;
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloOnnx::set_labels(const std::string& labels) {
|
||||
config_.labels = labels;
|
||||
}
|
||||
|
||||
const std::vector<DetectedObject>& YoloOnnx::objects() const {
|
||||
return detected_;
|
||||
}
|
||||
|
||||
std::string YoloOnnx::label_names() const {
|
||||
return config_.labels;
|
||||
}
|
||||
|
||||
} // namespace engines
|
||||
@@ -0,0 +1,49 @@
|
||||
#pragma once
|
||||
|
||||
#include "engines/engine.h"
|
||||
|
||||
#include <opencv2/core.hpp>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
namespace engines {
|
||||
|
||||
enum class YoloOutputLayout {
|
||||
AUTO, // detect from output tensor shape at load time
|
||||
NC_FIRST, // (1, 4+N, anchors) — YOLOv8, YOLOv11
|
||||
NC_LAST, // (1, anchors, 4+N) — rows-first, no objectness
|
||||
NC_LAST_OBJ // (1, anchors, 5+N) — YOLOv5, rows-first with objectness
|
||||
};
|
||||
|
||||
struct YoloOnnxConfig {
|
||||
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)
|
||||
std::string labels_file; // ZED_ONNX_LABELS_FILE (one class name per line)
|
||||
YoloOutputLayout output_layout = YoloOutputLayout::AUTO;
|
||||
|
||||
static YoloOnnxConfig from_env();
|
||||
};
|
||||
|
||||
class YoloOnnx : public Engine {
|
||||
public:
|
||||
explicit YoloOnnx(const YoloOnnxConfig& config);
|
||||
~YoloOnnx() 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
|
||||
void set_labels(const std::string& labels) override;
|
||||
const std::vector<DetectedObject>& objects() const override;
|
||||
std::string label_names() const override;
|
||||
|
||||
private:
|
||||
YoloOnnxConfig config_;
|
||||
struct Impl;
|
||||
std::unique_ptr<Impl> impl_;
|
||||
std::vector<DetectedObject> detected_;
|
||||
};
|
||||
|
||||
} // namespace engines
|
||||
@@ -0,0 +1,415 @@
|
||||
#include "engines/yoloworld_onnx.h"
|
||||
#include "engines/yolo_helpers.h"
|
||||
#include "util/clip_tokenizer.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <onnxruntime_cxx_api.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace engines {
|
||||
|
||||
using namespace detail;
|
||||
|
||||
// --- Impl ---
|
||||
|
||||
struct YoloWorldOnnx::Impl {
|
||||
// Visual model
|
||||
Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "yoloworld_onnx"};
|
||||
Ort::SessionOptions session_options;
|
||||
std::unique_ptr<Ort::Session> session;
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
|
||||
std::string input_name; // image input
|
||||
std::string text_input_name; // text_feats input
|
||||
std::string output_name;
|
||||
int num_classes = 0;
|
||||
|
||||
// Text encoder
|
||||
Ort::SessionOptions text_session_options;
|
||||
std::unique_ptr<Ort::Session> 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<float> cached_text_feats; // (num_classes * 512) flat, row-major
|
||||
std::string cached_labels;
|
||||
|
||||
std::vector<RawDetection> run_inference(const YoloWorldOnnxConfig& config,
|
||||
const cv::Mat& image_bgr);
|
||||
};
|
||||
|
||||
// --- Impl::run_inference ---
|
||||
|
||||
std::vector<RawDetection> YoloWorldOnnx::Impl::run_inference(
|
||||
const YoloWorldOnnxConfig& config, const cv::Mat& image_bgr)
|
||||
{
|
||||
if (!session) return {};
|
||||
if (cached_text_feats.empty()) return {};
|
||||
|
||||
int orig_w = image_bgr.cols;
|
||||
int orig_h = image_bgr.rows;
|
||||
|
||||
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<float> 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<cv::Vec3f>(h, w)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
std::vector<int64_t> image_shape{1, 3, sz, sz};
|
||||
Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info, input_data.data(), input_data.size(),
|
||||
image_shape.data(), image_shape.size());
|
||||
|
||||
std::vector<int64_t> text_shape{1, num_classes, 512};
|
||||
Ort::Value text_tensor = Ort::Value::CreateTensor<float>(
|
||||
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)};
|
||||
const char* output_name_cstr = output_name.c_str();
|
||||
|
||||
std::vector<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = session->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
input_names, input_tensors, 2,
|
||||
&output_name_cstr, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: inference failed: " << e.what() << std::endl;
|
||||
return {};
|
||||
}
|
||||
|
||||
// Output: (1, 4+num_classes, num_anchors) — NC_FIRST only
|
||||
auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
|
||||
if (out_shape.size() < 3) return {};
|
||||
|
||||
int num_anchors = static_cast<int>(out_shape[2]);
|
||||
int nc = num_classes;
|
||||
if (nc <= 0) return {};
|
||||
|
||||
const float* out = outputs[0].GetTensorData<float>();
|
||||
|
||||
std::vector<RawDetection> 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<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
|
||||
return nms(std::move(candidates), config.nms_threshold);
|
||||
}
|
||||
|
||||
// --- Config ---
|
||||
|
||||
YoloWorldOnnxConfig YoloWorldOnnxConfig::from_env() {
|
||||
YoloWorldOnnxConfig 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;
|
||||
}
|
||||
|
||||
// --- YoloWorldOnnx ---
|
||||
|
||||
YoloWorldOnnx::YoloWorldOnnx(const YoloWorldOnnxConfig& config)
|
||||
: config_(config), impl_(std::make_unique<Impl>()) {}
|
||||
|
||||
YoloWorldOnnx::~YoloWorldOnnx() = default;
|
||||
|
||||
void YoloWorldOnnx::enable(sl::Camera& /* camera */) {
|
||||
if (config_.model_path.empty()) {
|
||||
std::cerr << "YoloWorldOnnx: 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 << "YoloWorldOnnx: CUDA unavailable (" << e.what() << "), using CPU." << std::endl;
|
||||
}
|
||||
impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "YoloWorldOnnx: loading visual model (" << config_.model_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.model_path.c_str(), impl_->session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: failed to load visual model: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
size_t num_inputs = impl_->session->GetInputCount();
|
||||
if (num_inputs < 2) {
|
||||
std::cerr << "YoloWorldOnnx: model has only " << num_inputs << " input(s) — "
|
||||
"expected a two-input YOLO-World model. Use ZED_ENGINE=yolo_onnx for single-input models."
|
||||
<< std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->input_name = impl_->session->GetInputNameAllocated(0, impl_->allocator).get();
|
||||
impl_->text_input_name = impl_->session->GetInputNameAllocated(1, impl_->allocator).get();
|
||||
impl_->output_name = impl_->session->GetOutputNameAllocated(0, 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<int>(shape[1]) - 4;
|
||||
}
|
||||
|
||||
std::cout << "YoloWorldOnnx: visual model loaded."
|
||||
<< " Classes=" << impl_->num_classes
|
||||
<< std::endl;
|
||||
|
||||
// --- Load text encoder + tokenizer ---
|
||||
if (config_.text_encoder_path.empty()) {
|
||||
std::cerr << "YoloWorldOnnx: ZED_ONNX_TEXT_ENCODER_PATH is not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->text_session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "YoloWorldOnnx: loading text encoder (" << config_.text_encoder_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->text_session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.text_encoder_path.c_str(), impl_->text_session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: 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 << "YoloWorldOnnx: tokenizer loaded." << std::endl;
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: failed to load tokenizer: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
std::cerr << "YoloWorldOnnx: ZED_ONNX_CLIP_VOCAB / ZED_ONNX_CLIP_MERGES not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!config_.labels.empty()) {
|
||||
encode_labels(config_.labels);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::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<size_t>(N) * 512, 0.0f);
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
int encode_n = std::min(static_cast<int>(label_vec.size()), N);
|
||||
for (int i = 0; i < encode_n; i++) {
|
||||
auto tokens = impl_->tokenizer.tokenize(label_vec[i]);
|
||||
|
||||
std::vector<int64_t> input_shape{1, util::ClipTokenizer::SEQ_LEN};
|
||||
Ort::Value token_tensor = Ort::Value::CreateTensor<int64_t>(
|
||||
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<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = impl_->text_session->Run(
|
||||
Ort::RunOptions{nullptr}, &in, &token_tensor, 1, &out, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: text encoder failed for \"" << label_vec[i]
|
||||
<< "\": " << e.what() << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
const float* emb = outputs[0].GetTensorData<float>();
|
||||
|
||||
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 << "YoloWorldOnnx: encoded " << encode_n << " label(s): " << labels_str << std::endl;
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::detect(sl::Camera& camera) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
// Re-encode if labels changed (on-the-fly vocabulary switching)
|
||||
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;
|
||||
|
||||
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<int>(img_mat.getWidth());
|
||||
int orig_h = static_cast<int>(img_mat.getHeight());
|
||||
|
||||
cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
|
||||
cv::Mat bgr;
|
||||
cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR);
|
||||
|
||||
auto kept = impl_->run_inference(config_, bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
FusionResult fusion = fuse_bbox(pc_mat,
|
||||
static_cast<int>(det.x1), static_cast<int>(det.y1),
|
||||
static_cast<int>(det.x2), static_cast<int>(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);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::detect(const cv::Mat& image_bgr) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
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;
|
||||
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;
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::set_labels(const std::string& labels) {
|
||||
config_.labels = labels;
|
||||
encode_labels(labels);
|
||||
}
|
||||
|
||||
const std::vector<DetectedObject>& YoloWorldOnnx::objects() const {
|
||||
return detected_;
|
||||
}
|
||||
|
||||
std::string YoloWorldOnnx::label_names() const {
|
||||
return impl_->cached_labels.empty() ? config_.labels : impl_->cached_labels;
|
||||
}
|
||||
|
||||
} // namespace engines
|
||||
@@ -8,35 +8,35 @@
|
||||
|
||||
namespace engines {
|
||||
|
||||
struct OpenVocabConfig {
|
||||
struct YoloWorldOnnxConfig {
|
||||
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)
|
||||
// CLIP text encoder (required)
|
||||
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();
|
||||
static YoloWorldOnnxConfig from_env();
|
||||
};
|
||||
|
||||
class OpenVocab : public Engine {
|
||||
class YoloWorldOnnx : public Engine {
|
||||
public:
|
||||
explicit OpenVocab(const OpenVocabConfig& config);
|
||||
~OpenVocab() override;
|
||||
explicit YoloWorldOnnx(const YoloWorldOnnxConfig& config);
|
||||
~YoloWorldOnnx() 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
|
||||
void detect(const cv::Mat& image_bgr); // image-only; no 3D fusion
|
||||
void set_labels(const std::string& labels) override;
|
||||
const std::vector<DetectedObject>& objects() const override;
|
||||
std::string label_names() const override;
|
||||
|
||||
private:
|
||||
OpenVocabConfig config_;
|
||||
YoloWorldOnnxConfig config_;
|
||||
struct Impl;
|
||||
std::unique_ptr<Impl> impl_;
|
||||
std::vector<DetectedObject> detected_;
|
||||
+8
-4
@@ -1,5 +1,6 @@
|
||||
#include "engines/engine.h"
|
||||
#include "engines/open_vocab.h"
|
||||
#include "engines/yolo_onnx.h"
|
||||
#include "engines/yoloworld_onnx.h"
|
||||
#include "engines/zed_native.h"
|
||||
#include "util/env.h"
|
||||
#include "zed/camera.h"
|
||||
@@ -74,9 +75,12 @@ int main() {
|
||||
|
||||
std::string engine_name = util::get_env_str("ZED_ENGINE", "native");
|
||||
std::unique_ptr<engines::Engine> detector;
|
||||
if (engine_name == "open_vocab") {
|
||||
auto cfg = engines::OpenVocabConfig::from_env();
|
||||
detector = std::make_unique<engines::OpenVocab>(cfg);
|
||||
if (engine_name == "yolo_onnx" || engine_name == "onnx") {
|
||||
auto cfg = engines::YoloOnnxConfig::from_env();
|
||||
detector = std::make_unique<engines::YoloOnnx>(cfg);
|
||||
} else if (engine_name == "yoloworld_onnx" || engine_name == "open_vocab") {
|
||||
auto cfg = engines::YoloWorldOnnxConfig::from_env();
|
||||
detector = std::make_unique<engines::YoloWorldOnnx>(cfg);
|
||||
} else {
|
||||
auto cfg = engines::DetectionConfig::from_env();
|
||||
detector = std::make_unique<engines::ZedNative>(cfg);
|
||||
|
||||
@@ -1,24 +1,22 @@
|
||||
// Standalone test for the OpenVocab detection engine.
|
||||
// No dora framework, no live ZED camera — loads images from disk.
|
||||
// Standalone test for the YoloOnnx (closed-vocabulary) detection engine.
|
||||
// No dora framework, no live ZED camera — loads an image 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
|
||||
// bin/test_yolo_onnx [image_path]
|
||||
// image_path default: test/assets/yolo.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
|
||||
// Required env vars:
|
||||
// ZED_ONNX_MODEL_PATH path to YOLO ONNX model
|
||||
//
|
||||
// Optional env vars:
|
||||
// ZED_ONNX_LABELS comma-separated class names
|
||||
// ZED_ONNX_LABELS_FILE path to newline-separated class names file
|
||||
// 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_OUTPUT_FORMAT nc_first|nc_last|nc_last_obj|yolov5|yolov8|yolov11|auto
|
||||
|
||||
#include "engines/open_vocab.h"
|
||||
#include "engines/yolo_onnx.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
@@ -61,19 +59,28 @@ void draw_detections(cv::Mat& img,
|
||||
|
||||
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::putText(img, name + buf, {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) {
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
std::string image_path = (argc > 1) ? argv[1] : "test/assets/yolo.png";
|
||||
|
||||
auto config = engines::YoloOnnxConfig::from_env();
|
||||
|
||||
sl::Camera dummy_cam;
|
||||
engines::YoloOnnx engine(config);
|
||||
engine.enable(dummy_cam);
|
||||
|
||||
cv::Mat img = cv::imread(image_path);
|
||||
if (img.empty()) {
|
||||
std::cerr << "Error: cannot read image: " << image_path << std::endl;
|
||||
return false;
|
||||
return 1;
|
||||
}
|
||||
|
||||
engine.detect(img);
|
||||
@@ -101,42 +108,6 @@ bool run_test(engines::OpenVocab& engine, const std::string& image_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;
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
// Standalone test for the YoloWorldOnnx (open-vocabulary) detection engine.
|
||||
// No dora framework, no live ZED camera — loads an image from disk.
|
||||
//
|
||||
// Usage:
|
||||
// bin/test_yoloworld_onnx [image_path]
|
||||
// image_path default: test/assets/yoloworld.png
|
||||
//
|
||||
// Required env vars:
|
||||
// ZED_ONNX_MODEL_PATH path to YOLO-World ONNX model (two-input)
|
||||
// ZED_ONNX_TEXT_ENCODER_PATH path to clip_text_encoder.onnx
|
||||
// ZED_ONNX_CLIP_VOCAB path to clip_vocab.txt
|
||||
// ZED_ONNX_CLIP_MERGES path to clip_merges.txt
|
||||
// ZED_ONNX_LABELS comma-separated class labels
|
||||
//
|
||||
// Optional env vars:
|
||||
// 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)
|
||||
|
||||
#include "engines/yoloworld_onnx.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
namespace {
|
||||
|
||||
std::vector<std::string> parse_csv(const std::string& s) {
|
||||
std::vector<std::string> 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<engines::DetectedObject>& objects,
|
||||
const std::vector<std::string>& labels) {
|
||||
for (const auto& obj : objects) {
|
||||
int x1 = static_cast<int>(obj.bounding_box_2d[0]);
|
||||
int y1 = static_cast<int>(obj.bounding_box_2d[1]);
|
||||
int x2 = static_cast<int>(obj.bounding_box_2d[4]);
|
||||
int y2 = static_cast<int>(obj.bounding_box_2d[5]);
|
||||
|
||||
std::string name = (obj.label_id >= 0 && obj.label_id < static_cast<int>(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);
|
||||
|
||||
cv::rectangle(img, {x1, y1}, {x2, y2}, {0, 200, 0}, 2);
|
||||
cv::putText(img, name + buf, {x1, std::max(y1 - 6, 14)},
|
||||
cv::FONT_HERSHEY_SIMPLEX, 0.55, {0, 200, 0}, 2, cv::LINE_AA);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
std::string image_path = (argc > 1) ? argv[1] : "test/assets/yoloworld.png";
|
||||
|
||||
auto config = engines::YoloWorldOnnxConfig::from_env();
|
||||
|
||||
sl::Camera dummy_cam;
|
||||
engines::YoloWorldOnnx engine(config);
|
||||
engine.enable(dummy_cam);
|
||||
|
||||
cv::Mat img = cv::imread(image_path);
|
||||
if (img.empty()) {
|
||||
std::cerr << "Error: cannot read image: " << image_path << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
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<int>(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 0;
|
||||
}
|
||||
Reference in New Issue
Block a user