Add promptable YOLOE ONNX support

.cargo/config.toml

@@ -0,0 +1,2 @@
+[build]
+target-dir = "../yolo-rs-target"

.gitignore

@@ -8,6 +8,22 @@ devenv.local.nix
 # pre-commit
 .pre-commit-config.yaml
 
+# Editor settings
+.vscode/
+
+# Python local environment
+.venv/
+__pycache__/
+*.py[cod]
+
+# Local model weights
+/yoloe-26n-seg.pt
+/yoloe-v8s-seg.pt
+/yoloe-*.onnx
+/yoloe-*.onnx.data
+/mobileclip*.pt
+/mobileclip*.ts
+
 
 # Added by cargo
 

README.md

@@ -1,13 +1,13 @@
 # YOLO inference in Rust
 
-This project provides a Rust implementation of the YOLO v11 object detection model, enabling inference on images to identify objects along with their bounding boxes, labels, and confidence scores. It utilizes the YOLO v11 model in ONNX format and leverages the `ort` library for ONNX Runtime integration. The implementation is inspired by the [YOLOv8 example](https://github.com/pykeio/ort/tree/main/examples/yolov8) from the `ort` repository.
+This project provides a Rust implementation of YOLO-style ONNX object detection, enabling inference on images to identify objects along with their bounding boxes, labels, and confidence scores. It works out of the box with YOLO v8/v11 COCO exports and can also run exported closed-set YOLOE ONNX models when you provide the matching labels. The implementation is inspired by the [YOLOv8 example](https://github.com/pykeio/ort/tree/main/examples/yolov8) from the `ort` repository.
 
 See docs.rs for the [latest documentation](https://docs.rs/yolo-rs).
 
 ## Features
 
 - **Object Detection**: Detects objects within an image and provides their bounding boxes, labels, and confidence scores.
-- **ONNX Model Integration**: Employs the YOLO v11 model in ONNX format for efficient inference.
+- **YOLO-style ONNX Integration**: Supports default YOLO v8/v11 exports and exported closed-set YOLOE models.
 - **Rust Implementation**: Written entirely in Rust, ensuring performance and safety.
 - **ONNX Runtime**: Utilizes the `ort` library for executing the ONNX model.
 
@@ -31,6 +31,41 @@ On a MacBook Pro (2024) with M3 Max, it tooks about **57ms** to inferring an ima
 
 - [**yolo-cli**](examples/yolo-cli): Command-line interface for running YOLO inference on images.
 
+## YOLOE support
+
+This crate can run exported YOLOE ONNX models when they behave like regular fixed-class detectors. In practice that means exporting from Ultralytics after configuring the classes you want to detect, then passing the same labels into `yolo-rs`.
+
+Validated export flow in Python:
+
+```python
+from ultralytics import YOLOE
+
+model = YOLOE("yoloe-v8s-seg.pt")
+embeddings = model.get_text_pe(["person", "baseball bat", "baseball glove"])
+model.set_classes(["person", "baseball bat", "baseball glove"], embeddings)
+model.export(format="onnx")
+```
+
+That export produces a standard segmentation-style ONNX graph with:
+
+- input `images` shaped `[1, 3, 640, 640]`
+- output `output0` shaped `[1, 39, 8400]`
+- output `output1` shaped `[1, 32, 160, 160]`
+
+`yolo-rs` currently uses `output0` for box and class decoding and ignores `output1` and the trailing mask coefficients.
+
+Example CLI usage:
+
+```bash
+cargo run --release -p example-yolo-gui -- yoloe-v8s-seg.onnx examples/yolo-cli/data/baseball.jpg --labels-file target/yoloe_inspect_v839/yoloe.labels.txt --no-display
+```
+
+Current scope:
+
+- Closed-set YOLOE detection exported to ONNX is supported.
+- Detection heads with extra mask coefficients can be decoded for boxes and classes when you provide the label list.
+- Open-vocabulary text prompts, visual prompts, offline prompt embeddings, and segmentation mask decoding are not implemented in this crate yet.
+
 ## Acknowledgements
 
 This project is inspired by the [YOLOv8 example](https://github.com/pykeio/ort/tree/main/examples/yolov8) from the `ort` repository.

examples/yolo-cli/README.md

@@ -14,6 +14,24 @@ where:
     - You can export the model according to [Ultralytics' manual](https://docs.ultralytics.com/integrations/onnx/).
 - `data/baseball.jpg` is the path to the image file.
 
+For exported closed-set YOLOE models, pass the matching labels when you run the CLI:
+
+```bash
+cargo run --release exported-yoloe.onnx image.jpg --label person --label bus
+```
+
+For non-interactive validation, skip the GUI window:
+
+```bash
+cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt --no-display
+```
+
+You can also load labels from a file:
+
+```bash
+cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt
+```
+
 ## Downloading models
 
 ```shell

examples/yolo-cli/src/main.rs

@@ -1,3 +1,4 @@
+use std::fs;
 use std::path::PathBuf;
 
 use anyhow::{Context, Result};
@@ -13,11 +14,48 @@ struct Args {
     model_path: PathBuf,
     picture_path: PathBuf,
 
+    /// Override the class labels used to decode the model output.
+    ///
+    /// Pass this multiple times for exported YOLOE closed-set models, for example:
+    /// --label person --label bus
+    #[arg(long = "label")]
+    labels: Vec<String>,
+
+    /// Read class labels from a UTF-8 text file, one label per line.
+    #[arg(long)]
+    labels_file: Option<PathBuf>,
+
     #[arg(long)]
     probability_threshold: Option<f32>,
 
     #[arg(long)]
     iou_threshold: Option<f32>,
+
+    /// Skip opening the GUI window and only print detections.
+    #[arg(long)]
+    no_display: bool,
+}
+
+fn read_labels(args: &Args) -> Result<Option<Vec<String>>> {
+    let mut labels = if let Some(path) = &args.labels_file {
+        fs::read_to_string(path)
+            .with_context(|| format!("failed to read labels file {:?}", path.display()))?
+            .lines()
+            .map(str::trim)
+            .filter(|line| !line.is_empty())
+            .map(ToOwned::to_owned)
+            .collect::<Vec<_>>()
+    } else {
+        Vec::new()
+    };
+
+    labels.extend(args.labels.iter().cloned());
+
+    if labels.is_empty() {
+        Ok(None)
+    } else {
+        Ok(Some(labels))
+    }
 }
 
 #[show_image::main]
@@ -41,9 +79,17 @@ fn main() -> Result<()> {
         .with_context(|| format!("failed to open image {:?}", args.picture_path.display()))?;
 
     tracing::info!("Loading models {:?}…", args.model_path.display());
+    let labels = read_labels(&args)?;
     let mut model = {
-        let mut model = model::YoloModelSession::from_filename_v8(&args.model_path)
-            .with_context(|| format!("failed to load model {:?}", args.model_path.display()))?;
+        let mut model = if let Some(labels) = labels {
+            model::YoloModelSession::from_filename_with_labels(
+                &args.model_path,
+                labels.into_iter(),
+            )
+        } else {
+            model::YoloModelSession::from_filename_v8(&args.model_path)
+        }
+        .with_context(|| format!("failed to load model {:?}", args.model_path.display()))?;
 
         model.iou_threshold = args.iou_threshold;
         model.probability_threshold = args.probability_threshold;
@@ -117,6 +163,10 @@ fn main() -> Result<()> {
         );
     }
 
+    if args.no_display {
+        return Ok(());
+    }
+
     let overlay: show_image::Image = dt.into();
 
     tracing::info!("Displaying image…");

src/error.rs

@@ -10,6 +10,22 @@ pub enum YoloError {
     OrtInputError(ort::Error),
     #[error("run inference: {0}")]
     OrtInferenceError(ort::Error),
-    #[error("extract sensor: {0}")]
-    OrtExtractSensorError(ort::Error),
+    #[error("extract tensor: {0}")]
+    OrtExtractTensorError(ort::Error),
+    #[error("model has no inputs")]
+    MissingModelInput,
+    #[error("model has no outputs")]
+    MissingModelOutput,
+    #[error("unsupported model output shape {0:?}; expected a 3D detection tensor")]
+    InvalidOutputShape(Vec<usize>),
+    #[error(
+        "model output has {available} detection channels, but {required} are required for 4 box coordinates plus {label_count} labels"
+    )]
+    InsufficientDetectionChannels {
+        available: usize,
+        required: usize,
+        label_count: usize,
+    },
+    #[error("detected class index {class_id} is out of range for {label_count} labels")]
+    UnknownClassId { class_id: usize, label_count: usize },
 }

src/lib.rs

@@ -10,7 +10,7 @@ use arcstr::ArcStr;
 use error::YoloError;
 use image::{DynamicImage, GenericImageView, Rgba, imageops::FilterType};
 use model::YoloModelSession;
-use ndarray::{Array4, ArrayBase, ArrayView4, Axis, s};
+use ndarray::{Array4, ArrayBase, ArrayView3, ArrayView4};
 use ort::{inputs, value::TensorRef};
 
 #[derive(Debug, Clone, Copy)]
@@ -87,7 +87,8 @@ pub fn image_to_yolo_input_tensor(original_image: &DynamicImage) -> YoloInput {
 /// Inference on the YOLO model, returning the detected entities.
 ///
 /// The input tensor should be obtained from the [`image_to_yolo_input_tensor`] function.
-/// The [`YoloModelSession`] can be obtained from the [`YoloModelSession::from_filename_v8`] method.
+/// The [`YoloModelSession`] can be obtained from the [`YoloModelSession::from_filename_v8`]
+/// or [`YoloModelSession::from_filename_with_labels`] methods.
 pub fn inference(
     model: &mut YoloModelSession,
     YoloInputView {
@@ -96,6 +97,35 @@ pub fn inference(
         raw_height,
     }: YoloInputView,
 ) -> Result<Vec<YoloEntityOutput>, YoloError> {
+    fn output_layout(
+        output: ArrayView3<'_, f32>,
+        label_count: usize,
+    ) -> Result<(usize, usize), YoloError> {
+        if output.shape()[0] != 1 {
+            return Err(YoloError::InvalidOutputShape(output.shape().to_vec()));
+        }
+
+        let channels_axis = if output.shape()[1] <= output.shape()[2] {
+            1
+        } else {
+            2
+        };
+        let detection_axis = if channels_axis == 1 { 2 } else { 1 };
+
+        let channel_count = output.shape()[channels_axis];
+        let required_channels = 4 + label_count;
+
+        if channel_count < required_channels {
+            return Err(YoloError::InsufficientDetectionChannels {
+                available: channel_count,
+                required: required_channels,
+                label_count,
+            });
+        }
+
+        Ok((channels_axis, output.shape()[detection_axis]))
+    }
+
     fn intersection(box1: &BoundingBox, box2: &BoundingBox) -> f32 {
         (box1.x2.min(box2.x2) - box1.x1.max(box2.x1))
             * (box1.y2.min(box2.y2) - box1.y1.max(box2.y1))
@@ -143,32 +173,50 @@ pub fn inference(
     let iou_threshold = model.get_iou_threshold();
     let probability_threshold = model.get_probability_threshold();
     let labels = model.get_labels().to_vec();
+    let input_name = model.get_input_name()?.to_owned();
+    let output_name = model.get_output_name()?.to_owned();
 
-    // Run YOLOv8 inference
-    let inputs = inputs!["images" => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?];
+    // Run YOLO-style inference using the session's configured input/output names.
+    let inputs = inputs![input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?];
     let outputs = model
         .as_mut()
         .run(inputs)
         .map_err(YoloError::OrtInferenceError)?;
-    let output = outputs["output0"]
+    let output = outputs[output_name.as_str()]
         .try_extract_array::<f32>()
-        .map_err(YoloError::OrtExtractSensorError)?
-        .reversed_axes();
-    let output = output.slice(s![.., .., 0]);
+        .map_err(YoloError::OrtExtractTensorError)?;
+    let output = output
+        .view()
+        .into_dimensionality::<ndarray::Ix3>()
+        .map_err(|_| YoloError::InvalidOutputShape(output.shape().to_vec()))?;
+    let (channels_axis, detection_count) = output_layout(output, labels.len())?;
 
     // Turn the output tensor into bounding boxes
-    let boxes = output
-        .axis_iter(Axis(0))
+    let boxes = (0..detection_count)
         .filter_map(|row| {
+            let row = if channels_axis == 1 {
+                output.slice(ndarray::s![0, .., row])
+            } else {
+                output.slice(ndarray::s![0, row, ..])
+            };
+
             let (class_id, prob) = row
                 .iter()
-                .skip(4) // skip bounding box coordinates
+                .skip(4)
+                .take(labels.len())
                 .enumerate()
                 .map(|(index, value)| (index, *value))
                 .reduce(|accum, row| if row.1 > accum.1 { row } else { accum })
                 .filter(|(_, prob)| *prob >= probability_threshold)?;
 
-            let label = labels[class_id].clone();
+            let label = labels
+                .get(class_id)
+                .cloned()
+                .ok_or(YoloError::UnknownClassId {
+                    class_id,
+                    label_count: labels.len(),
+                })
+                .ok()?;
 
             let xc = row[0_usize] / 640. * (raw_width as f32);
             let yc = row[1_usize] / 640. * (raw_height as f32);