From 6ecde63c805720ab428a40ff1e45b77a3fe03b09 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 09:58:00 +0200 Subject: [PATCH 01/10] Add configurable YOLO and YOLOE model support --- README.md | 30 ++++++++++++++- examples/yolo-cli/README.md | 12 ++++++ examples/yolo-cli/src/main.rs | 46 +++++++++++++++++++++- src/error.rs | 20 +++++++++- src/lib.rs | 72 +++++++++++++++++++++++++++++------ src/model.rs | 53 ++++++++++++++++++++++++++ 6 files changed, 215 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 936e736..69fc754 100644 --- a/README.md +++ b/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,32 @@ 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`. + +Example export flow in Python: + +```python +from ultralytics import YOLOE + +model = YOLOE("yoloe-26s-seg.pt") +model.set_classes(["person", "bus"]) +model.export(format="onnx") +``` + +Example CLI usage: + +```bash +cargo run --release -p example-yolo-gui -- exported-yoloe.onnx image.jpg --label person --label bus +``` + +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. diff --git a/examples/yolo-cli/README.md b/examples/yolo-cli/README.md index a953e25..7416086 100644 --- a/examples/yolo-cli/README.md +++ b/examples/yolo-cli/README.md @@ -14,6 +14,18 @@ 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 +``` + +You can also load labels from a file: + +```bash +cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt +``` + ## Downloading models ```shell diff --git a/examples/yolo-cli/src/main.rs b/examples/yolo-cli/src/main.rs index 87c0519..05e2df3 100644 --- a/examples/yolo-cli/src/main.rs +++ b/examples/yolo-cli/src/main.rs @@ -1,3 +1,4 @@ +use std::fs; use std::path::PathBuf; use anyhow::{Context, Result}; @@ -13,6 +14,17 @@ 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, + + /// Read class labels from a UTF-8 text file, one label per line. + #[arg(long)] + labels_file: Option, + #[arg(long)] probability_threshold: Option, @@ -20,6 +32,28 @@ struct Args { iou_threshold: Option, } +fn read_labels(args: &Args) -> Result>> { + 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::>() + } else { + Vec::new() + }; + + labels.extend(args.labels.iter().cloned()); + + if labels.is_empty() { + Ok(None) + } else { + Ok(Some(labels)) + } +} + #[show_image::main] fn main() -> Result<()> { tracing_subscriber::fmt::init(); @@ -41,9 +75,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; diff --git a/src/error.rs b/src/error.rs index c8fac32..6f01316 100644 --- a/src/error.rs +++ b/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), + #[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 }, } diff --git a/src/lib.rs b/src/lib.rs index 1d3bd6a..49b2db7 100644 --- a/src/lib.rs +++ b/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, 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::() - .map_err(YoloError::OrtExtractSensorError)? - .reversed_axes(); - let output = output.slice(s![.., .., 0]); + .map_err(YoloError::OrtExtractTensorError)?; + let output = output + .view() + .into_dimensionality::() + .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); diff --git a/src/model.rs b/src/model.rs index 362456c..fb5bdd7 100644 --- a/src/model.rs +++ b/src/model.rs @@ -13,6 +13,8 @@ use crate::error::YoloError; pub struct YoloModelSession { pub session: ort::session::Session, pub labels: Vec, + pub input_name: Option, + pub output_name: Option, pub probability_threshold: Option, // default = 0.5 pub iou_threshold: Option, // default = 0.7 @@ -29,6 +31,25 @@ impl YoloModelSession { Self { session, labels: labels.map(Into::into).collect(), + input_name: None, + output_name: None, + probability_threshold: None, + iou_threshold: None, + } + } + + /// Wrap an ONNX session to a [`YoloModelSession`] with explicit input/output names. + pub fn new_with_io_names( + session: ort::session::Session, + labels: impl Iterator>, + input_name: impl Into, + output_name: impl Into, + ) -> Self { + Self { + session, + labels: labels.map(Into::into).collect(), + input_name: Some(input_name.into()), + output_name: Some(output_name.into()), probability_threshold: None, iou_threshold: None, } @@ -122,6 +143,8 @@ impl YoloModelSession { Self { session, labels: LABELS.to_vec(), + input_name: None, + output_name: None, probability_threshold: None, iou_threshold: None, } @@ -143,6 +166,22 @@ impl YoloModelSession { Ok(Self::new_v8(session)) } + /// Load a YOLO-style ONNX model from a filename with caller-provided labels. + /// + /// This is useful for exported YOLOE closed-set models where the class list is not + /// the default COCO vocabulary. + pub fn from_filename_with_labels( + filename: impl AsRef, + labels: impl Iterator>, + ) -> Result { + let session = ort::session::Session::builder() + .map_err(YoloError::OrtSessionBuildError)? + .commit_from_file(filename) + .map_err(YoloError::OrtSessionLoadError)?; + + Ok(Self::new(session, labels)) + } + pub fn get_labels(&self) -> &[ArcStr] { &self.labels } @@ -154,6 +193,20 @@ impl YoloModelSession { pub fn get_iou_threshold(&self) -> f32 { self.iou_threshold.unwrap_or(0.7) } + + pub fn get_input_name(&self) -> Result<&str, YoloError> { + self.input_name + .as_deref() + .or_else(|| self.session.inputs().first().map(|input| input.name())) + .ok_or(YoloError::MissingModelInput) + } + + pub fn get_output_name(&self) -> Result<&str, YoloError> { + self.output_name + .as_deref() + .or_else(|| self.session.outputs().first().map(|output| output.name())) + .ok_or(YoloError::MissingModelOutput) + } } impl AsRef for YoloModelSession { From 61a60faae678ca46ba59e2b22a1c3b07e7f43981 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:00:27 +0200 Subject: [PATCH 02/10] Add headless YOLOE validation flow --- README.md | 17 +++++++++++++---- examples/yolo-cli/README.md | 6 ++++++ examples/yolo-cli/src/main.rs | 8 ++++++++ 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 69fc754..389f3e4 100644 --- a/README.md +++ b/README.md @@ -35,20 +35,29 @@ On a MacBook Pro (2024) with M3 Max, it tooks about **57ms** to inferring an ima 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`. -Example export flow in Python: +Validated export flow in Python: ```python from ultralytics import YOLOE -model = YOLOE("yoloe-26s-seg.pt") -model.set_classes(["person", "bus"]) +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 -- exported-yoloe.onnx image.jpg --label person --label bus +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: diff --git a/examples/yolo-cli/README.md b/examples/yolo-cli/README.md index 7416086..a3568ef 100644 --- a/examples/yolo-cli/README.md +++ b/examples/yolo-cli/README.md @@ -20,6 +20,12 @@ For exported closed-set YOLOE models, pass the matching labels when you run the 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 diff --git a/examples/yolo-cli/src/main.rs b/examples/yolo-cli/src/main.rs index 05e2df3..1295a63 100644 --- a/examples/yolo-cli/src/main.rs +++ b/examples/yolo-cli/src/main.rs @@ -30,6 +30,10 @@ struct Args { #[arg(long)] iou_threshold: Option, + + /// Skip opening the GUI window and only print detections. + #[arg(long)] + no_display: bool, } fn read_labels(args: &Args) -> Result>> { @@ -159,6 +163,10 @@ fn main() -> Result<()> { ); } + if args.no_display { + return Ok(()); + } + let overlay: show_image::Image = dt.into(); tracing::info!("Displaying image…"); From d9e2310be93e0ef5abe1bc47b84b9151bd96576b Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:12:55 +0200 Subject: [PATCH 03/10] Ignore local YOLOE export artifacts --- .gitignore | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/.gitignore b/.gitignore index b972a7f..fb94426 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,19 @@ devenv.local.nix # pre-commit .pre-commit-config.yaml +# 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 From 698e2aa028eb830a8e87638e103ea894e3844c50 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:17:14 +0200 Subject: [PATCH 04/10] Hide and clean local generated artifacts --- .vscode/settings.json | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 .vscode/settings.json diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..22c5c2f --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,17 @@ +{ + "files.exclude": { + "**/.venv": true, + "**/.mypy_cache": true, + "**/target": true, + "mobileclip*.pt": true, + "mobileclip*.ts": true, + "yoloe-*.onnx": true, + "yoloe-*.onnx.data": true, + "yoloe-*.pt": true + }, + "search.exclude": { + "**/.venv": true, + "**/.mypy_cache": true, + "**/target": true + } +} \ No newline at end of file From e4233de0ff7216eeecfbe3d392cf16efc07643c1 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:22:56 +0200 Subject: [PATCH 05/10] Move Cargo build output outside workspace --- .cargo/config.toml | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 .cargo/config.toml diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..58ec2d6 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[build] +target-dir = "../yolo-rs-target" \ No newline at end of file From 9ed56d15d1a7101d85cdacbb0b453f753e10be52 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:37:31 +0200 Subject: [PATCH 06/10] Ignore local VS Code settings --- .gitignore | 3 +++ .vscode/settings.json | 17 ----------------- 2 files changed, 3 insertions(+), 17 deletions(-) delete mode 100644 .vscode/settings.json diff --git a/.gitignore b/.gitignore index fb94426..d5d31e5 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,9 @@ devenv.local.nix # pre-commit .pre-commit-config.yaml +# Editor settings +.vscode/ + # Python local environment .venv/ __pycache__/ diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index 22c5c2f..0000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "files.exclude": { - "**/.venv": true, - "**/.mypy_cache": true, - "**/target": true, - "mobileclip*.pt": true, - "mobileclip*.ts": true, - "yoloe-*.onnx": true, - "yoloe-*.onnx.data": true, - "yoloe-*.pt": true - }, - "search.exclude": { - "**/.venv": true, - "**/.mypy_cache": true, - "**/target": true - } -} \ No newline at end of file From ddfae5ce998b46c5ebf713ae11a6ca7017e8ea3d Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:38:33 +0200 Subject: [PATCH 07/10] Ignore local Cargo configuration --- .cargo/config.toml | 2 -- .gitignore | 3 +++ 2 files changed, 3 insertions(+), 2 deletions(-) delete mode 100644 .cargo/config.toml diff --git a/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index 58ec2d6..0000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,2 +0,0 @@ -[build] -target-dir = "../yolo-rs-target" \ No newline at end of file diff --git a/.gitignore b/.gitignore index d5d31e5..5cd4521 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,9 @@ devenv.local.nix # Editor settings .vscode/ +# Local Cargo config +.cargo/ + # Python local environment .venv/ __pycache__/ From b436b9bb7527b88259ef0f6f4a9b07dc31eb546a Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:54:07 +0200 Subject: [PATCH 08/10] Decode segmentation masks from YOLOE ONNX --- README.md | 4 +- examples/yolo-cli/README.md | 4 + examples/yolo-cli/src/main.rs | 34 ++- src/error.rs | 12 ++ src/lib.rs | 375 +++++++++++++++++++++++----------- src/model.rs | 8 + 6 files changed, 317 insertions(+), 120 deletions(-) diff --git a/README.md b/README.md index 389f3e4..57e3018 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ That export produces a standard segmentation-style ONNX graph with: - 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. +`yolo-rs` now also decodes `output1` and the trailing mask coefficients when you run a segmentation-style export through the example CLI. Example CLI usage: @@ -64,7 +65,8 @@ 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. +- Segmentation-style exports can decode instance masks from `output1` and the mask coefficients stored in `output0`. +- Runtime open-vocabulary text prompts and visual prompts are not implemented for ONNX in this crate. Ultralytics' exported ONNX graphs expose only the `images` input, so prompts must be fused into the model before export via `set_classes(...)`. ## Acknowledgements diff --git a/examples/yolo-cli/README.md b/examples/yolo-cli/README.md index a3568ef..237df05 100644 --- a/examples/yolo-cli/README.md +++ b/examples/yolo-cli/README.md @@ -26,6 +26,10 @@ For non-interactive validation, skip the GUI window: cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt --no-display ``` +For segmentation-style exports, the CLI will also decode and draw instance masks when the model exposes a second ONNX output with mask prototypes. + +Open-vocabulary runtime prompting is not available for exported ONNX models in this crate. For YOLOE, prompts must be configured before export so the exported graph behaves like a fixed-class detector. + You can also load labels from a file: ```bash diff --git a/examples/yolo-cli/src/main.rs b/examples/yolo-cli/src/main.rs index 1295a63..a584695 100644 --- a/examples/yolo-cli/src/main.rs +++ b/examples/yolo-cli/src/main.rs @@ -6,7 +6,7 @@ use clap::Parser; use ort::execution_providers::{CUDAExecutionProvider, CoreMLExecutionProvider}; use raqote::{DrawOptions, DrawTarget, LineJoin, PathBuilder, SolidSource, Source, StrokeStyle}; use show_image::{AsImageView, WindowOptions, event}; -use yolo_rs::{YoloEntityOutput, image_to_yolo_input_tensor, inference, model}; +use yolo_rs::{YoloEntityOutput, YoloSegmentationOutput, image_to_yolo_input_tensor, inference, inference_segment, model}; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -104,13 +104,39 @@ fn main() -> Result<()> { tracing::info!("Running inference…"); let now = std::time::Instant::now(); - let result = inference(&mut model, input.view())?; + let has_mask_output = model.session.outputs().len() > 1; + let result = if has_mask_output { + let (img_width, img_height) = (original_img.width(), original_img.height()); + let mut dt = DrawTarget::new(img_width as _, img_height as _); + let result = inference_segment(&mut model, input.view())?; + let data = dt.get_data_mut(); + for YoloSegmentationOutput { entity, mask } in &result { + let fill = match entity.label.as_str() { + "baseball bat" => 0x40001080u32, + "baseball glove" => 0x40208040u32, + _ => 0x40801040u32, + }; + + for (x, y, pixel) in mask.enumerate_pixels() { + if pixel.0[0] > 0 { + let index = (y as usize) * (img_width as usize) + (x as usize); + data[index] = fill; + } + } + } + + (result.into_iter().map(|output| output.entity).collect::>(), dt) + } else { + let (img_width, img_height) = (original_img.width(), original_img.height()); + let dt = DrawTarget::new(img_width as _, img_height as _); + (inference(&mut model, input.view())?, dt) + }; + tracing::info!("Inference took {:?}", now.elapsed()); tracing::debug!("Drawing bounding boxes…"); let (img_width, img_height) = (original_img.width(), original_img.height()); - - let mut dt = DrawTarget::new(img_width as _, img_height as _); + let (result, mut dt) = result; for YoloEntityOutput { bounding_box: bbox, diff --git a/src/error.rs b/src/error.rs index 6f01316..c90db0c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -16,8 +16,12 @@ pub enum YoloError { MissingModelInput, #[error("model has no outputs")] MissingModelOutput, + #[error("model has no segmentation mask output")] + MissingMaskOutput, #[error("unsupported model output shape {0:?}; expected a 3D detection tensor")] InvalidOutputShape(Vec), + #[error("unsupported mask output shape {0:?}; expected a 4D mask prototype tensor")] + InvalidMaskShape(Vec), #[error( "model output has {available} detection channels, but {required} are required for 4 box coordinates plus {label_count} labels" )] @@ -26,6 +30,14 @@ pub enum YoloError { required: usize, label_count: usize, }, + #[error( + "model output has {available} mask coefficients, but {required} are required for {prototype_count} mask prototypes" + )] + InsufficientMaskCoefficients { + available: usize, + required: usize, + prototype_count: usize, + }, #[error("detected class index {class_id} is out of range for {label_count} labels")] UnknownClassId { class_id: usize, label_count: usize }, } diff --git a/src/lib.rs b/src/lib.rs index 49b2db7..dece6c0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,7 +8,7 @@ pub mod model; use arcstr::ArcStr; use error::YoloError; -use image::{DynamicImage, GenericImageView, Rgba, imageops::FilterType}; +use image::{DynamicImage, GenericImageView, GrayImage, Luma, Rgba, imageops::FilterType}; use model::YoloModelSession; use ndarray::{Array4, ArrayBase, ArrayView3, ArrayView4}; use ort::{inputs, value::TensorRef}; @@ -57,6 +57,196 @@ pub struct YoloEntityOutput { pub confidence: f32, } +#[derive(Debug, Clone)] +pub struct YoloSegmentationOutput { + pub entity: YoloEntityOutput, + pub mask: GrayImage, +} + +#[derive(Debug, Clone)] +struct DecodedDetection { + entity: YoloEntityOutput, + mask_coefficients: Vec, +} + +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)) +} + +fn union(box1: &BoundingBox, box2: &BoundingBox) -> f32 { + ((box1.x2 - box1.x1) * (box1.y2 - box1.y1)) + ((box2.x2 - box2.x1) * (box2.y2 - box2.y1)) + - intersection(box1, box2) +} + +fn non_maximum_suppression( + mut boxes: Vec, + iou_threshold: f32, +) -> Vec { + if boxes.is_empty() { + return Vec::new(); + } + + boxes.sort_unstable_by(|a, b| b.entity.confidence.total_cmp(&a.entity.confidence)); + + let mut result = Vec::with_capacity(boxes.len()); + for current in boxes { + if result.iter().all(|selected: &DecodedDetection| { + let iou = intersection(&selected.entity.bounding_box, ¤t.entity.bounding_box) + / union(&selected.entity.bounding_box, ¤t.entity.bounding_box); + iou < iou_threshold + }) { + result.push(current); + } + } + + result.shrink_to_fit(); + result +} + +fn decode_detections( + output: ArrayView3<'_, f32>, + labels: &[ArcStr], + probability_threshold: f32, + raw_width: u32, + raw_height: u32, + expected_mask_coefficients: Option, +) -> Result, YoloError> { + let (channels_axis, detection_count) = output_layout(output, labels.len())?; + + Ok((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) + .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 + .get(class_id) + .cloned() + .ok_or(YoloError::UnknownClassId { + class_id, + label_count: labels.len(), + }) + .ok()?; + + let mask_coefficients = row + .iter() + .skip(4 + labels.len()) + .copied() + .collect::>(); + + if let Some(required) = expected_mask_coefficients + && mask_coefficients.len() < required + { + return Some(Err(YoloError::InsufficientMaskCoefficients { + available: mask_coefficients.len(), + required, + prototype_count: required, + })); + } + + let xc = row[0_usize] / 640. * (raw_width as f32); + let yc = row[1_usize] / 640. * (raw_height as f32); + let w = row[2_usize] / 640. * (raw_width as f32); + let h = row[3_usize] / 640. * (raw_height as f32); + + Some(Ok(DecodedDetection { + entity: YoloEntityOutput { + bounding_box: BoundingBox { + x1: xc - w / 2., + y1: yc - h / 2., + x2: xc + w / 2., + y2: yc + h / 2., + }, + label, + confidence: prob, + }, + mask_coefficients, + })) + }) + .collect::, _>>()?) +} + +fn decode_segmentation_masks( + detections: Vec, + prototypes: ArrayView3<'_, f32>, + raw_width: u32, + raw_height: u32, +) -> Result, YoloError> { + let prototype_count = prototypes.shape()[0]; + let mask_height = prototypes.shape()[1]; + let mask_width = prototypes.shape()[2]; + + let mut results = Vec::with_capacity(detections.len()); + for detection in detections { + let mut mask = GrayImage::new(raw_width, raw_height); + let bbox = detection.entity.bounding_box; + + let start_x = bbox.x1.floor().max(0.0) as u32; + let end_x = bbox.x2.ceil().min(raw_width as f32) as u32; + let start_y = bbox.y1.floor().max(0.0) as u32; + let end_y = bbox.y2.ceil().min(raw_height as f32) as u32; + + for y in start_y..end_y { + let proto_y = ((y as usize) * mask_height) / (raw_height as usize); + for x in start_x..end_x { + let proto_x = ((x as usize) * mask_width) / (raw_width as usize); + + let value = detection + .mask_coefficients + .iter() + .take(prototype_count) + .enumerate() + .map(|(index, coefficient)| coefficient * prototypes[[index, proto_y, proto_x]]) + .sum::(); + + if value > 0.0 { + mask.put_pixel(x, y, Luma([255])); + } + } + } + + results.push(YoloSegmentationOutput { + entity: detection.entity, + mask, + }); + } + + Ok(results) +} + /// Convert an image to a YOLO input tensor. /// /// The input image is resized to 640x640 and normalized to the range [0, 1]. @@ -97,77 +287,6 @@ pub fn inference( raw_height, }: YoloInputView, ) -> Result, 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)) - } - - fn union(box1: &BoundingBox, box2: &BoundingBox) -> f32 { - ((box1.x2 - box1.x1) * (box1.y2 - box1.y1)) + ((box2.x2 - box2.x1) * (box2.y2 - box2.y1)) - - intersection(box1, box2) - } - - fn non_maximum_suppression( - mut boxes: Vec, - iou_threshold: f32, - ) -> Vec { - // Early return if no boxes are provided - if boxes.is_empty() { - return Vec::new(); - } - - // Sort boxes by confidence descending using sort_unstable_by for better performance - boxes.sort_unstable_by(|a, b| b.confidence.total_cmp(&a.confidence)); - - let mut result = Vec::with_capacity(boxes.len()); - - // Iterate through each box and select it if it doesn't overlap significantly with already selected boxes - for current in boxes.into_iter() { - // Check if the current box has a high IoU with any box in the result - // Using `iter().all()` ensures we short-circuit on the first overlap found - if result.iter().all(|selected: &YoloEntityOutput| { - let iou = intersection(&selected.bounding_box, ¤t.bounding_box) - / union(&selected.bounding_box, ¤t.bounding_box); - iou < iou_threshold - }) { - result.push(current); - } - } - - result.shrink_to_fit(); - - result - } - // Due to the lifetime of the model, we need to clone the // labels and thresholds early. let iou_threshold = model.get_iou_threshold(); @@ -189,53 +308,79 @@ pub fn inference( .view() .into_dimensionality::() .map_err(|_| YoloError::InvalidOutputShape(output.shape().to_vec()))?; - let (channels_axis, detection_count) = output_layout(output, labels.len())?; + let boxes = decode_detections( + output, + &labels, + probability_threshold, + raw_width, + raw_height, + None, + )?; - // Turn the output tensor into bounding boxes - 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, ..]) - }; + Ok(non_maximum_suppression(boxes, iou_threshold) + .into_iter() + .map(|decoded| decoded.entity) + .collect()) +} - let (class_id, prob) = row - .iter() - .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)?; +/// Inference on a segmentation-style YOLO model, returning the detected entities and decoded masks. +/// +/// This is intended for YOLO segmentation exports, including exported closed-set YOLOE segmentation models. +pub fn inference_segment( + model: &mut YoloModelSession, + YoloInputView { + tensor_view, + raw_width, + raw_height, + }: YoloInputView, +) -> Result, YoloError> { + 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(); + let mask_output_name = model.get_mask_output_name()?.to_owned(); - let label = labels - .get(class_id) - .cloned() - .ok_or(YoloError::UnknownClassId { - class_id, - label_count: labels.len(), - }) - .ok()?; + 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 xc = row[0_usize] / 640. * (raw_width as f32); - let yc = row[1_usize] / 640. * (raw_height as f32); - let w = row[2_usize] / 640. * (raw_width as f32); - let h = row[3_usize] / 640. * (raw_height as f32); + let detection_output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let detection_output = detection_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(detection_output.shape().to_vec()))?; - Some(YoloEntityOutput { - bounding_box: BoundingBox { - x1: xc - w / 2., - y1: yc - h / 2., - x2: xc + w / 2., - y2: yc + h / 2., - }, - label, - confidence: prob, - }) - }) - .collect::>(); + let mask_output = outputs[mask_output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let mask_output = mask_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidMaskShape(mask_output.shape().to_vec()))?; + if mask_output.shape()[0] != 1 { + return Err(YoloError::InvalidMaskShape(mask_output.shape().to_vec())); + } - // Perform non-maximum suppression (NMS) - Ok(non_maximum_suppression(boxes, iou_threshold)) + let prototype_count = mask_output.shape()[1]; + let boxes = decode_detections( + detection_output, + &labels, + probability_threshold, + raw_width, + raw_height, + Some(prototype_count), + )?; + let boxes = non_maximum_suppression(boxes, iou_threshold); + + decode_segmentation_masks( + boxes, + mask_output.slice(ndarray::s![0, .., .., ..]), + raw_width, + raw_height, + ) } diff --git a/src/model.rs b/src/model.rs index fb5bdd7..4f37678 100644 --- a/src/model.rs +++ b/src/model.rs @@ -207,6 +207,14 @@ impl YoloModelSession { .or_else(|| self.session.outputs().first().map(|output| output.name())) .ok_or(YoloError::MissingModelOutput) } + + pub fn get_mask_output_name(&self) -> Result<&str, YoloError> { + self.session + .outputs() + .get(1) + .map(|output| output.name()) + .ok_or(YoloError::MissingMaskOutput) + } } impl AsRef for YoloModelSession { From 9b0bd562b772c55bb79d61a2023d231c11d9acb7 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 11:59:49 +0200 Subject: [PATCH 09/10] Align YOLOE decoding with upstream flow --- README.md | 8 +-- src/lib.rs | 180 ++++++++++++++++++++++++++++++++++++++++++----------- 2 files changed, 147 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 57e3018..3ce3ba8 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # YOLO inference in Rust -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. +This project provides a Rust implementation of YOLO-style ONNX inference, enabling detection and segmentation on images with bounding boxes, labels, confidence scores, and decoded masks. 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. +- **Object Detection and Segmentation**: Detects objects within an image and provides their bounding boxes, labels, confidence scores, and segmentation masks when the model exports them. - **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. @@ -52,8 +52,7 @@ That export produces a standard segmentation-style ONNX graph with: - 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. -`yolo-rs` now also decodes `output1` and the trailing mask coefficients when you run a segmentation-style export through the example CLI. +`yolo-rs` follows the same broad exported ONNX contract as Ultralytics' YOLO segmentation examples: it letterboxes the input image, decodes boxes from `output0`, and reconstructs instance masks from the trailing mask coefficients in `output0` plus the prototype tensor in `output1`. Example CLI usage: @@ -66,6 +65,7 @@ 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. - Segmentation-style exports can decode instance masks from `output1` and the mask coefficients stored in `output0`. +- Preprocessing and box rescaling now follow YOLO-style letterboxing instead of stretching the image to `640x640`. - Runtime open-vocabulary text prompts and visual prompts are not implemented for ONNX in this crate. Ultralytics' exported ONNX graphs expose only the `images` input, so prompts must be fused into the model before export via `set_classes(...)`. ## Acknowledgements diff --git a/src/lib.rs b/src/lib.rs index dece6c0..e1d21fc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ -//! A Rust library for the YOLO v11 object detection model. +//! A Rust library for YOLO-style ONNX detection and segmentation models. //! -//! This library provides a high-level API for running the YOLO v11 object detection model. -//! Currently, it supports only the inference. +//! This library provides a high-level API for running YOLO-style ONNX models, +//! including exported closed-set YOLOE segmentation models. pub mod error; pub mod model; @@ -13,6 +13,8 @@ use model::YoloModelSession; use ndarray::{Array4, ArrayBase, ArrayView3, ArrayView4}; use ort::{inputs, value::TensorRef}; +const DEFAULT_MODEL_SIZE: u32 = 640; + #[derive(Debug, Clone, Copy)] pub struct BoundingBox { pub x1: f32, @@ -26,6 +28,11 @@ pub struct YoloInput { pub tensor: Array4, // 640x640 pub raw_width: u32, pub raw_height: u32, + pub model_width: u32, + pub model_height: u32, + pub resize_ratio: f32, + pub pad_w: f32, + pub pad_h: f32, } impl YoloInput { @@ -34,6 +41,11 @@ impl YoloInput { tensor_view: self.tensor.view(), raw_width: self.raw_width, raw_height: self.raw_height, + model_width: self.model_width, + model_height: self.model_height, + resize_ratio: self.resize_ratio, + pad_w: self.pad_w, + pad_h: self.pad_h, } } } @@ -43,6 +55,11 @@ pub struct YoloInputView<'a> { pub tensor_view: ArrayView4<'a, f32>, pub raw_width: u32, pub raw_height: u32, + pub model_width: u32, + pub model_height: u32, + pub resize_ratio: f32, + pub pad_w: f32, + pub pad_h: f32, } #[derive(Debug, Clone)] @@ -69,6 +86,17 @@ struct DecodedDetection { mask_coefficients: Vec, } +#[derive(Debug, Clone, Copy)] +struct LetterboxInfo { + raw_width: u32, + raw_height: u32, + model_width: u32, + model_height: u32, + resize_ratio: f32, + pad_w: f32, + pad_h: f32, +} + 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())); @@ -129,8 +157,7 @@ fn decode_detections( output: ArrayView3<'_, f32>, labels: &[ArcStr], probability_threshold: f32, - raw_width: u32, - raw_height: u32, + letterbox: LetterboxInfo, expected_mask_coefficients: Option, ) -> Result, YoloError> { let (channels_axis, detection_count) = output_layout(output, labels.len())?; @@ -177,18 +204,18 @@ fn decode_detections( })); } - let xc = row[0_usize] / 640. * (raw_width as f32); - let yc = row[1_usize] / 640. * (raw_height as f32); - let w = row[2_usize] / 640. * (raw_width as f32); - let h = row[3_usize] / 640. * (raw_height as f32); + let x1 = ((row[0_usize] - row[2_usize] / 2.0) - letterbox.pad_w) / letterbox.resize_ratio; + let y1 = ((row[1_usize] - row[3_usize] / 2.0) - letterbox.pad_h) / letterbox.resize_ratio; + let x2 = ((row[0_usize] + row[2_usize] / 2.0) - letterbox.pad_w) / letterbox.resize_ratio; + let y2 = ((row[1_usize] + row[3_usize] / 2.0) - letterbox.pad_h) / letterbox.resize_ratio; Some(Ok(DecodedDetection { entity: YoloEntityOutput { bounding_box: BoundingBox { - x1: xc - w / 2., - y1: yc - h / 2., - x2: xc + w / 2., - y2: yc + h / 2., + x1: x1.clamp(0.0, letterbox.raw_width as f32), + y1: y1.clamp(0.0, letterbox.raw_height as f32), + x2: x2.clamp(0.0, letterbox.raw_width as f32), + y2: y2.clamp(0.0, letterbox.raw_height as f32), }, label, confidence: prob, @@ -202,28 +229,30 @@ fn decode_detections( fn decode_segmentation_masks( detections: Vec, prototypes: ArrayView3<'_, f32>, - raw_width: u32, - raw_height: u32, + letterbox: LetterboxInfo, ) -> Result, YoloError> { let prototype_count = prototypes.shape()[0]; let mask_height = prototypes.shape()[1]; let mask_width = prototypes.shape()[2]; + let mask_pad_w = letterbox.pad_w * (mask_width as f32 / letterbox.model_width as f32); + let mask_pad_h = letterbox.pad_h * (mask_height as f32 / letterbox.model_height as f32); + let mask_content_width = (mask_width as f32 - (mask_pad_w * 2.0)).max(1.0); + let mask_content_height = (mask_height as f32 - (mask_pad_h * 2.0)).max(1.0); + let mut results = Vec::with_capacity(detections.len()); for detection in detections { - let mut mask = GrayImage::new(raw_width, raw_height); + let mut mask = GrayImage::new(letterbox.raw_width, letterbox.raw_height); let bbox = detection.entity.bounding_box; let start_x = bbox.x1.floor().max(0.0) as u32; - let end_x = bbox.x2.ceil().min(raw_width as f32) as u32; + let end_x = bbox.x2.ceil().min(letterbox.raw_width as f32) as u32; let start_y = bbox.y1.floor().max(0.0) as u32; - let end_y = bbox.y2.ceil().min(raw_height as f32) as u32; - - for y in start_y..end_y { - let proto_y = ((y as usize) * mask_height) / (raw_height as usize); - for x in start_x..end_x { - let proto_x = ((x as usize) * mask_width) / (raw_width as usize); + let end_y = bbox.y2.ceil().min(letterbox.raw_height as f32) as u32; + let mut proto_mask = vec![0.0f32; mask_width * mask_height]; + for proto_y in 0..mask_height { + for proto_x in 0..mask_width { let value = detection .mask_coefficients .iter() @@ -231,8 +260,17 @@ fn decode_segmentation_masks( .enumerate() .map(|(index, coefficient)| coefficient * prototypes[[index, proto_y, proto_x]]) .sum::(); + proto_mask[proto_y * mask_width + proto_x] = value; + } + } + + for y in start_y..end_y { + for x in start_x..end_x { + let scaled_x = ((x as f32 + 0.5) / letterbox.raw_width as f32) * mask_content_width + mask_pad_w; + let scaled_y = ((y as f32 + 0.5) / letterbox.raw_height as f32) * mask_content_height + mask_pad_h; + let value = bilinear_sample(&proto_mask, mask_width, mask_height, scaled_x, scaled_y); - if value > 0.0 { + if value > 0.5 { mask.put_pixel(x, y, Luma([255])); } } @@ -247,6 +285,28 @@ fn decode_segmentation_masks( Ok(results) } +fn bilinear_sample(buffer: &[f32], width: usize, height: usize, x: f32, y: f32) -> f32 { + let x = x.clamp(0.0, (width.saturating_sub(1)) as f32); + let y = y.clamp(0.0, (height.saturating_sub(1)) as f32); + + let x0 = x.floor() as usize; + let y0 = y.floor() as usize; + let x1 = (x0 + 1).min(width.saturating_sub(1)); + let y1 = (y0 + 1).min(height.saturating_sub(1)); + + let dx = x - x0 as f32; + let dy = y - y0 as f32; + + let top_left = buffer[y0 * width + x0]; + let top_right = buffer[y0 * width + x1]; + let bottom_left = buffer[y1 * width + x0]; + let bottom_right = buffer[y1 * width + x1]; + + let top = top_left + dx * (top_right - top_left); + let bottom = bottom_left + dx * (bottom_right - bottom_left); + top + dy * (bottom - top) +} + /// Convert an image to a YOLO input tensor. /// /// The input image is resized to 640x640 and normalized to the range [0, 1]. @@ -255,12 +315,28 @@ fn decode_segmentation_masks( /// You can pass the resulting tensor to the [`inference`] function. /// Note that you might need to call [`YoloInput::view`] to get a view of the tensor. pub fn image_to_yolo_input_tensor(original_image: &DynamicImage) -> YoloInput { - let mut input = ArrayBase::zeros((1, 3, 640, 640)); - - let image = original_image.resize_exact(640, 640, FilterType::CatmullRom); + let model_width = DEFAULT_MODEL_SIZE; + let model_height = DEFAULT_MODEL_SIZE; + let mut input = ArrayBase::from_elem( + (1, 3, model_height as usize, model_width as usize), + 114.0f32 / 255.0, + ); + + let raw_width = original_image.width(); + let raw_height = original_image.height(); + let resize_ratio = (model_width as f32 / raw_width as f32) + .min(model_height as f32 / raw_height as f32); + let resized_width = ((raw_width as f32) * resize_ratio).round() as u32; + let resized_height = ((raw_height as f32) * resize_ratio).round() as u32; + let pad_w = (model_width as f32 - resized_width as f32) / 2.0; + let pad_h = (model_height as f32 - resized_height as f32) / 2.0; + let left = (pad_w - 0.1).round().max(0.0) as usize; + let top = (pad_h - 0.1).round().max(0.0) as usize; + + let image = original_image.resize_exact(resized_width, resized_height, FilterType::CatmullRom); for (x, y, Rgba([r, g, b, _])) in image.pixels() { - let x = x as usize; - let y = y as usize; + let x = left + x as usize; + let y = top + y as usize; input[[0, 0, y, x]] = (r as f32) / 255.; input[[0, 1, y, x]] = (g as f32) / 255.; @@ -269,8 +345,13 @@ pub fn image_to_yolo_input_tensor(original_image: &DynamicImage) -> YoloInput { YoloInput { tensor: input, - raw_width: original_image.width(), - raw_height: original_image.height(), + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, } } @@ -285,6 +366,11 @@ pub fn inference( tensor_view, raw_width, raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, }: YoloInputView, ) -> Result, YoloError> { // Due to the lifetime of the model, we need to clone the @@ -308,12 +394,20 @@ pub fn inference( .view() .into_dimensionality::() .map_err(|_| YoloError::InvalidOutputShape(output.shape().to_vec()))?; + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; let boxes = decode_detections( output, &labels, probability_threshold, - raw_width, - raw_height, + letterbox, None, )?; @@ -332,6 +426,11 @@ pub fn inference_segment( tensor_view, raw_width, raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, }: YoloInputView, ) -> Result, YoloError> { let iou_threshold = model.get_iou_threshold(); @@ -367,12 +466,20 @@ pub fn inference_segment( } let prototype_count = mask_output.shape()[1]; + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; let boxes = decode_detections( detection_output, &labels, probability_threshold, - raw_width, - raw_height, + letterbox, Some(prototype_count), )?; let boxes = non_maximum_suppression(boxes, iou_threshold); @@ -380,7 +487,6 @@ pub fn inference_segment( decode_segmentation_masks( boxes, mask_output.slice(ndarray::s![0, .., .., ..]), - raw_width, - raw_height, + letterbox, ) } From 57437456563f3e39dcc1672ee7478062cff1a497 Mon Sep 17 00:00:00 2001 From: vlordier Date: Fri, 17 Apr 2026 12:08:23 +0200 Subject: [PATCH 10/10] Add runtime YOLOE prompt support --- Cargo.lock | 389 ++++++++++++++++++++++++++++- Cargo.toml | 1 + README.md | 33 ++- examples/yolo-cli/README.md | 14 +- examples/yolo-cli/src/main.rs | 64 ++++- scripts/export_promptable_yoloe.py | 121 +++++++++ src/error.rs | 16 ++ src/lib.rs | 155 ++++++++++++ src/model.rs | 18 ++ src/prompt.rs | 205 +++++++++++++++ 10 files changed, 1006 insertions(+), 10 deletions(-) create mode 100644 scripts/export_promptable_yoloe.py create mode 100644 src/prompt.rs diff --git a/Cargo.lock b/Cargo.lock index af222d4..b7fbbd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -40,11 +40,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", + "getrandom 0.3.4", "once_cell", + "serde", "version_check", "zerocopy", ] +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "aligned" version = "0.4.3" @@ -246,7 +257,7 @@ dependencies = [ "anyhow", "arrayvec", "log", - "nom", + "nom 8.0.0", "num-rational", "v_frame", ] @@ -275,6 +286,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.22.1" @@ -404,6 +421,15 @@ dependencies = [ "vec_map", ] +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.2.56" @@ -475,7 +501,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3538270d33cc669650c4b093848450d380def10c331d38c768e34cac80576e6e" dependencies = [ "termcolor", - "unicode-width", + "unicode-width 0.1.14", ] [[package]] @@ -496,6 +522,33 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf43edc576402991846b093a7ca18a3477e0ef9c588cde84964b5d3e43016642" +[[package]] +name = "compact_str" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "console" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width 0.2.2", + "windows-sys 0.61.2", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -606,6 +659,50 @@ dependencies = [ "winapi", ] +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.117", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "der" version = "0.7.10" @@ -616,6 +713,37 @@ dependencies = [ "zeroize", ] +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.117", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -643,6 +771,12 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "equator" version = "0.4.2" @@ -679,6 +813,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "euclid" version = "0.22.13" @@ -775,6 +918,12 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.1.5" @@ -1107,6 +1256,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "image" version = "0.25.9" @@ -1169,6 +1324,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" +dependencies = [ + "console", + "portable-atomic", + "unicode-width 0.2.2", + "unit-prefix", + "web-time", +] + [[package]] name = "instant" version = "0.1.13" @@ -1368,6 +1536,22 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69" +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + [[package]] name = "malloc_buf" version = "0.0.6" @@ -1436,6 +1620,12 @@ dependencies = [ "paste", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -1458,6 +1648,28 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "moxcms" version = "0.7.11" @@ -1581,6 +1793,16 @@ dependencies = [ "memoffset", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nom" version = "8.0.0" @@ -1773,6 +1995,28 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "onig" +version = "6.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0" +dependencies = [ + "bitflags 2.11.0", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f86c6eef3d6df15f23bcfb6af487cbd2fed4e5581d58d5bf1f5f8b7f6727dc" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "openssl" version = "0.10.75" @@ -2188,6 +2432,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -2225,6 +2480,35 @@ dependencies = [ "bitflags 2.11.0", ] +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + [[package]] name = "renderdoc-sys" version = "1.1.0" @@ -2286,6 +2570,12 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "schannel" version = "0.1.28" @@ -2356,6 +2646,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -2510,6 +2801,18 @@ dependencies = [ "num-traits", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom 7.1.3", + "serde", + "unicode-segmentation", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -2672,6 +2975,40 @@ dependencies = [ "strict-num", ] +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "indicatif", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "toml_datetime" version = "0.6.11" @@ -2764,25 +3101,58 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" + [[package]] name = "unicode-width" version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode-xid" version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + [[package]] name = "ureq" version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdc97a28575b85cfedf2a7e7d3cc64b3e11bd8ac766666318003abbacc7a21fc" dependencies = [ - "base64", + "base64 0.22.1", "der", "log", "native-tls", @@ -2800,7 +3170,7 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d81f9efa9df032be5934a46a068815a10a042b494b6a58cb0a1a97bb5467ed6f" dependencies = [ - "base64", + "base64 0.22.1", "http", "httparse", "log", @@ -3051,6 +3421,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki-root-certs" version = "1.0.6" @@ -3528,6 +3908,7 @@ dependencies = [ "ndarray", "ort", "thiserror 2.0.18", + "tokenizers", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 5465913..3a67a82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ image = "0.25.9" ndarray = "0.17" ort = "2.0.0-rc.11" thiserror = "2.0.18" +tokenizers = "0.22.1" [workspace] members = [ diff --git a/README.md b/README.md index 3ce3ba8..2dd1f22 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,38 @@ Current scope: - Detection heads with extra mask coefficients can be decoded for boxes and classes when you provide the label list. - Segmentation-style exports can decode instance masks from `output1` and the mask coefficients stored in `output0`. - Preprocessing and box rescaling now follow YOLO-style letterboxing instead of stretching the image to `640x640`. -- Runtime open-vocabulary text prompts and visual prompts are not implemented for ONNX in this crate. Ultralytics' exported ONNX graphs expose only the `images` input, so prompts must be fused into the model before export via `set_classes(...)`. +- Runtime open-vocabulary prompting is available only for custom promptable YOLOE ONNX exports that expose a prompt embedding input. Ultralytics' stock ONNX exporter still emits single-input graphs, so those default exports must keep using `set_classes(...)` before export. + +## True Runtime Prompting + +`yolo-rs` now also supports promptable YOLOE ONNX graphs with a second model input for prompt embeddings. + +That path is different from Ultralytics' stock `model.export(format="onnx")`, which still emits a single-input ONNX and bakes prompts in ahead of time. For true runtime open-vocabulary prompting you need three artifacts: + +- a promptable YOLOE detector ONNX with inputs `images` and `prompt_embeddings` +- a prompt-encoder ONNX that converts tokenized text into the detector's prompt embedding space +- a tokenizer JSON compatible with that prompt encoder + +This repository includes a helper export script at [scripts/export_promptable_yoloe.py](/Users/vincent/Work/yolo-rs/scripts/export_promptable_yoloe.py) to generate those assets from upstream YOLOE weights. + +Example CLI usage once those assets exist: + +```bash +cargo run --release -p example-yolo-gui -- \ + exported-promptable-yoloe/yoloe-promptable.onnx \ + examples/yolo-cli/data/baseball.jpg \ + --prompt person \ + --prompt "baseball bat" \ + --prompt "baseball glove" \ + --prompt-encoder-model exported-promptable-yoloe/yoloe-prompt-encoder.onnx \ + --prompt-tokenizer exported-promptable-yoloe/tokenizer.json +``` + +Current runtime-prompt scope: + +- `yolo-rs` can run promptable two-input YOLOE ONNX graphs with runtime text prompts. +- The prompt encoder is a separate ONNX model; prompt strings are tokenized in Rust and encoded at runtime. +- Stock Ultralytics ONNX exports remain single-input closed-set graphs, so they still require `set_classes(...)` before export. ## Acknowledgements diff --git a/examples/yolo-cli/README.md b/examples/yolo-cli/README.md index 237df05..fde0a2d 100644 --- a/examples/yolo-cli/README.md +++ b/examples/yolo-cli/README.md @@ -28,7 +28,19 @@ cargo run --release exported-yoloe.onnx image.jpg --labels-file labels.txt --no- For segmentation-style exports, the CLI will also decode and draw instance masks when the model exposes a second ONNX output with mask prototypes. -Open-vocabulary runtime prompting is not available for exported ONNX models in this crate. For YOLOE, prompts must be configured before export so the exported graph behaves like a fixed-class detector. +Open-vocabulary runtime prompting is available only for custom promptable YOLOE ONNX exports. Ultralytics' default ONNX export still produces a single-input fixed-class graph. + +For true runtime prompting with a promptable two-input YOLOE ONNX graph, pass prompts together with a prompt encoder and tokenizer: + +```bash +cargo run --release promptable-yoloe.onnx image.jpg \ + --prompt person \ + --prompt bus \ + --prompt-encoder-model yoloe-prompt-encoder.onnx \ + --prompt-tokenizer tokenizer.json +``` + +This path expects a custom promptable export rather than Ultralytics' default single-input ONNX export. See [scripts/export_promptable_yoloe.py](/Users/vincent/Work/yolo-rs/scripts/export_promptable_yoloe.py) for a helper script. You can also load labels from a file: diff --git a/examples/yolo-cli/src/main.rs b/examples/yolo-cli/src/main.rs index a584695..6028574 100644 --- a/examples/yolo-cli/src/main.rs +++ b/examples/yolo-cli/src/main.rs @@ -6,7 +6,11 @@ use clap::Parser; use ort::execution_providers::{CUDAExecutionProvider, CoreMLExecutionProvider}; use raqote::{DrawOptions, DrawTarget, LineJoin, PathBuilder, SolidSource, Source, StrokeStyle}; use show_image::{AsImageView, WindowOptions, event}; -use yolo_rs::{YoloEntityOutput, YoloSegmentationOutput, image_to_yolo_input_tensor, inference, inference_segment, model}; +use yolo_rs::{ + YoloEntityOutput, YoloSegmentationOutput, image_to_yolo_input_tensor, inference, + inference_segment, inference_segment_with_prompts, inference_with_prompts, model, + prompt::YoloPromptEncoderSession, +}; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -21,6 +25,18 @@ struct Args { #[arg(long = "label")] labels: Vec, + /// Runtime open-vocabulary prompts for a promptable YOLOE ONNX model. + #[arg(long = "prompt")] + prompts: Vec, + + /// Path to a prompt-encoder ONNX model that converts tokenized text to prompt embeddings. + #[arg(long)] + prompt_encoder_model: Option, + + /// Path to a tokenizer JSON compatible with the prompt encoder. + #[arg(long)] + prompt_tokenizer: Option, + /// Read class labels from a UTF-8 text file, one label per line. #[arg(long)] labels_file: Option, @@ -80,8 +96,11 @@ fn main() -> Result<()> { tracing::info!("Loading models {:?}…", args.model_path.display()); let labels = read_labels(&args)?; + let prompt_mode = !args.prompts.is_empty(); let mut model = { - let mut model = if let Some(labels) = labels { + let mut model = if prompt_mode { + model::YoloModelSession::from_filename(&args.model_path) + } else if let Some(labels) = labels { model::YoloModelSession::from_filename_with_labels( &args.model_path, labels.into_iter(), @@ -100,6 +119,32 @@ fn main() -> Result<()> { tracing::debug!("Converting image to tensor…"); let input = image_to_yolo_input_tensor(&original_img); + let prompt_embeddings = if prompt_mode { + let prompt_encoder_model = args + .prompt_encoder_model + .as_ref() + .context("--prompt-encoder-model is required when using --prompt")?; + let prompt_tokenizer = args + .prompt_tokenizer + .as_ref() + .context("--prompt-tokenizer is required when using --prompt")?; + tracing::info!("Encoding runtime prompts…"); + let mut prompt_encoder = YoloPromptEncoderSession::from_files( + prompt_encoder_model, + prompt_tokenizer, + ) + .with_context(|| { + format!( + "failed to load prompt encoder {:?} and tokenizer {:?}", + prompt_encoder_model.display(), + prompt_tokenizer.display() + ) + })?; + Some(prompt_encoder.encode(args.prompts.iter().map(String::as_str))?) + } else { + None + }; + // Run YOLOv8 inference tracing::info!("Running inference…"); @@ -108,7 +153,11 @@ fn main() -> Result<()> { let result = if has_mask_output { let (img_width, img_height) = (original_img.width(), original_img.height()); let mut dt = DrawTarget::new(img_width as _, img_height as _); - let result = inference_segment(&mut model, input.view())?; + let result = if let Some(prompt_embeddings) = prompt_embeddings.as_ref() { + inference_segment_with_prompts(&mut model, input.view(), prompt_embeddings.view())? + } else { + inference_segment(&mut model, input.view())? + }; let data = dt.get_data_mut(); for YoloSegmentationOutput { entity, mask } in &result { let fill = match entity.label.as_str() { @@ -129,7 +178,14 @@ fn main() -> Result<()> { } else { let (img_width, img_height) = (original_img.width(), original_img.height()); let dt = DrawTarget::new(img_width as _, img_height as _); - (inference(&mut model, input.view())?, dt) + ( + if let Some(prompt_embeddings) = prompt_embeddings.as_ref() { + inference_with_prompts(&mut model, input.view(), prompt_embeddings.view())? + } else { + inference(&mut model, input.view())? + }, + dt, + ) }; tracing::info!("Inference took {:?}", now.elapsed()); diff --git a/scripts/export_promptable_yoloe.py b/scripts/export_promptable_yoloe.py new file mode 100644 index 0000000..dd1aaa0 --- /dev/null +++ b/scripts/export_promptable_yoloe.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from ultralytics import YOLOE +from ultralytics.nn.modules.head import Detect +from ultralytics.nn.text_model import build_text_model + + +class PromptableDetector(torch.nn.Module): + def __init__(self, checkpoint: str) -> None: + super().__init__() + wrapper = YOLOE(checkpoint) + self.model = wrapper.model + self.model.eval().float() + + for module in self.model.modules(): + if isinstance(module, Detect): + module.dynamic = False + module.export = True + module.format = "onnx" + module.max_det = 300 + + def forward(self, images: torch.Tensor, prompt_embeddings: torch.Tensor): + return self.model.predict(images, tpe=prompt_embeddings) + + +class PromptEncoder(torch.nn.Module): + def __init__(self, checkpoint: str) -> None: + super().__init__() + wrapper = YOLOE(checkpoint) + self.model = wrapper.model + self.model.eval().float() + self.text_model = build_text_model(self.model.args.get("text_model"), device=torch.device("cpu")) + self.text_model.eval() + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + text_features = self.text_model.encode_text(input_ids, dtype=torch.float32) + text_features = text_features.reshape(1, input_ids.shape[0], text_features.shape[-1]) + head = self.model.model[-1] + return head.get_tpe(text_features) + + +def save_tokenizer(encoder: PromptEncoder, output_path: Path) -> None: + tokenizer = encoder.text_model.tokenizer + output_path.parent.mkdir(parents=True, exist_ok=True) + + if hasattr(tokenizer, "save"): + tokenizer.save(str(output_path)) + return + + if hasattr(tokenizer, "save_pretrained"): + tokenizer.save_pretrained(str(output_path.parent)) + return + + raise RuntimeError("Tokenizer does not support save() or save_pretrained(); export it manually.") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", help="Path to a YOLOE checkpoint such as yoloe-v8s-seg.pt") + parser.add_argument("--output-dir", type=Path, default=Path("exported-promptable-yoloe")) + parser.add_argument("--opset", type=int, default=17) + parser.add_argument("--height", type=int, default=640) + parser.add_argument("--width", type=int, default=640) + args = parser.parse_args() + + args.output_dir.mkdir(parents=True, exist_ok=True) + + detector = PromptableDetector(args.checkpoint) + prompt_encoder = PromptEncoder(args.checkpoint) + + sample_image = torch.zeros((1, 3, args.height, args.width), dtype=torch.float32) + sample_tokens = prompt_encoder.text_model.tokenize(["object"]).to(torch.long) + sample_prompt_embeddings = prompt_encoder(sample_tokens) + + detector_path = args.output_dir / "yoloe-promptable.onnx" + prompt_encoder_path = args.output_dir / "yoloe-prompt-encoder.onnx" + tokenizer_path = args.output_dir / "tokenizer.json" + + detector_output_names = ["output0", "output1"] if "seg" in Path(args.checkpoint).stem else ["output0"] + + torch.onnx.export( + detector, + (sample_image, sample_prompt_embeddings), + detector_path, + opset_version=args.opset, + input_names=["images", "prompt_embeddings"], + output_names=detector_output_names, + dynamic_axes={ + "images": {0: "batch"}, + "prompt_embeddings": {1: "prompt_count"}, + "output0": {0: "batch", 2: "anchors"}, + **({"output1": {0: "batch", 2: "mask_height", 3: "mask_width"}} if len(detector_output_names) > 1 else {}), + }, + ) + + torch.onnx.export( + prompt_encoder, + sample_tokens, + prompt_encoder_path, + opset_version=args.opset, + input_names=["input_ids"], + output_names=["prompt_embeddings"], + dynamic_axes={ + "input_ids": {0: "prompt_count", 1: "token_count"}, + "prompt_embeddings": {1: "prompt_count"}, + }, + ) + + save_tokenizer(prompt_encoder, tokenizer_path) + + print(f"Wrote detector ONNX to {detector_path}") + print(f"Wrote prompt encoder ONNX to {prompt_encoder_path}") + print(f"Wrote tokenizer to {tokenizer_path}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/error.rs b/src/error.rs index c90db0c..fd35fe8 100644 --- a/src/error.rs +++ b/src/error.rs @@ -16,12 +16,28 @@ pub enum YoloError { MissingModelInput, #[error("model has no outputs")] MissingModelOutput, + #[error("model has no prompt embedding input")] + MissingPromptInput, #[error("model has no segmentation mask output")] MissingMaskOutput, #[error("unsupported model output shape {0:?}; expected a 3D detection tensor")] InvalidOutputShape(Vec), #[error("unsupported mask output shape {0:?}; expected a 4D mask prototype tensor")] InvalidMaskShape(Vec), + #[error("unsupported prompt embedding shape {0:?}; expected [1, prompts, embedding_dim]")] + InvalidPromptEmbeddingShape(Vec), + #[error("unsupported prompt encoder output shape {0:?}; expected [prompts, embedding_dim] or [1, prompts, embedding_dim]")] + InvalidPromptEncoderOutputShape(Vec), + #[error("prompt encoder model exposes {0} inputs, but only 1 to 3 text-encoder inputs are supported")] + UnsupportedPromptEncoderInputCount(usize), + #[error("prompt list is empty")] + EmptyPromptList, + #[error("prompt label count {labels} does not match prompt embedding count {embeddings}")] + PromptLabelEmbeddingMismatch { labels: usize, embeddings: usize }, + #[error("load tokenizer: {0}")] + TokenizerLoadError(String), + #[error("tokenize prompts: {0}")] + PromptTokenizationError(String), #[error( "model output has {available} detection channels, but {required} are required for 4 box coordinates plus {label_count} labels" )] diff --git a/src/lib.rs b/src/lib.rs index e1d21fc..763c26b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ pub mod error; pub mod model; +pub mod prompt; use arcstr::ArcStr; use error::YoloError; @@ -12,6 +13,7 @@ use image::{DynamicImage, GenericImageView, GrayImage, Luma, Rgba, imageops::Fil use model::YoloModelSession; use ndarray::{Array4, ArrayBase, ArrayView3, ArrayView4}; use ort::{inputs, value::TensorRef}; +use prompt::YoloPromptEmbeddingsView; const DEFAULT_MODEL_SIZE: u32 = 640; @@ -80,6 +82,24 @@ pub struct YoloSegmentationOutput { pub mask: GrayImage, } +fn validate_prompt_embeddings( + prompts: YoloPromptEmbeddingsView<'_>, +) -> Result, YoloError> { + if prompts.embeddings.shape().len() != 3 || prompts.embeddings.shape()[0] != 1 { + return Err(YoloError::InvalidPromptEmbeddingShape( + prompts.embeddings.shape().to_vec(), + )); + } + if prompts.labels.len() != prompts.embeddings.shape()[1] { + return Err(YoloError::PromptLabelEmbeddingMismatch { + labels: prompts.labels.len(), + embeddings: prompts.embeddings.shape()[1], + }); + } + + Ok(prompts) +} + #[derive(Debug, Clone)] struct DecodedDetection { entity: YoloEntityOutput, @@ -417,6 +437,65 @@ pub fn inference( .collect()) } +pub fn inference_with_prompts( + model: &mut YoloModelSession, + YoloInputView { + tensor_view, + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }: YoloInputView, + prompts: YoloPromptEmbeddingsView<'_>, +) -> Result, YoloError> { + let prompts = validate_prompt_embeddings(prompts)?; + let iou_threshold = model.get_iou_threshold(); + let probability_threshold = model.get_probability_threshold(); + let input_name = model.get_input_name()?.to_owned(); + let prompt_input_name = model.get_prompt_input_name()?.to_owned(); + let output_name = model.get_output_name()?.to_owned(); + + let inputs = inputs![ + input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?, + prompt_input_name.as_str() => TensorRef::from_array_view(prompts.embeddings).map_err(YoloError::OrtInputError)?, + ]; + let outputs = model + .as_mut() + .run(inputs) + .map_err(YoloError::OrtInferenceError)?; + let output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let output = output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(output.shape().to_vec()))?; + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; + let boxes = decode_detections( + output, + prompts.labels, + probability_threshold, + letterbox, + None, + )?; + + Ok(non_maximum_suppression(boxes, iou_threshold) + .into_iter() + .map(|decoded| decoded.entity) + .collect()) +} + /// Inference on a segmentation-style YOLO model, returning the detected entities and decoded masks. /// /// This is intended for YOLO segmentation exports, including exported closed-set YOLOE segmentation models. @@ -490,3 +569,79 @@ pub fn inference_segment( letterbox, ) } + +pub fn inference_segment_with_prompts( + model: &mut YoloModelSession, + YoloInputView { + tensor_view, + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }: YoloInputView, + prompts: YoloPromptEmbeddingsView<'_>, +) -> Result, YoloError> { + let prompts = validate_prompt_embeddings(prompts)?; + let iou_threshold = model.get_iou_threshold(); + let probability_threshold = model.get_probability_threshold(); + let input_name = model.get_input_name()?.to_owned(); + let prompt_input_name = model.get_prompt_input_name()?.to_owned(); + let output_name = model.get_output_name()?.to_owned(); + let mask_output_name = model.get_mask_output_name()?.to_owned(); + + let inputs = inputs![ + input_name.as_str() => TensorRef::from_array_view(tensor_view).map_err(YoloError::OrtInputError)?, + prompt_input_name.as_str() => TensorRef::from_array_view(prompts.embeddings).map_err(YoloError::OrtInputError)?, + ]; + let outputs = model + .as_mut() + .run(inputs) + .map_err(YoloError::OrtInferenceError)?; + + let detection_output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let detection_output = detection_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidOutputShape(detection_output.shape().to_vec()))?; + + let mask_output = outputs[mask_output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + let mask_output = mask_output + .view() + .into_dimensionality::() + .map_err(|_| YoloError::InvalidMaskShape(mask_output.shape().to_vec()))?; + if mask_output.shape()[0] != 1 { + return Err(YoloError::InvalidMaskShape(mask_output.shape().to_vec())); + } + + let letterbox = LetterboxInfo { + raw_width, + raw_height, + model_width, + model_height, + resize_ratio, + pad_w, + pad_h, + }; + let prototype_count = mask_output.shape()[1]; + let boxes = decode_detections( + detection_output, + prompts.labels, + probability_threshold, + letterbox, + Some(prototype_count), + )?; + let boxes = non_maximum_suppression(boxes, iou_threshold); + + decode_segmentation_masks( + boxes, + mask_output.slice(ndarray::s![0, .., .., ..]), + letterbox, + ) +} diff --git a/src/model.rs b/src/model.rs index 4f37678..69988b3 100644 --- a/src/model.rs +++ b/src/model.rs @@ -21,6 +21,16 @@ pub struct YoloModelSession { } impl YoloModelSession { + /// Wrap a generic ONNX session to a [`YoloModelSession`] without predefined labels. + pub fn from_filename(filename: impl AsRef) -> Result { + let session = ort::session::Session::builder() + .map_err(YoloError::OrtSessionBuildError)? + .commit_from_file(filename) + .map_err(YoloError::OrtSessionLoadError)?; + + Ok(Self::new(session, std::iter::empty::())) + } + /// Wrap a ONNX session to a [`YoloModelSession`]. /// /// The `session` is the ONNX runtime session, and the `labels` are the YOLO labels. @@ -208,6 +218,14 @@ impl YoloModelSession { .ok_or(YoloError::MissingModelOutput) } + pub fn get_prompt_input_name(&self) -> Result<&str, YoloError> { + self.session + .inputs() + .get(1) + .map(|input| input.name()) + .ok_or(YoloError::MissingPromptInput) + } + pub fn get_mask_output_name(&self) -> Result<&str, YoloError> { self.session .outputs() diff --git a/src/prompt.rs b/src/prompt.rs new file mode 100644 index 0000000..3f32cf6 --- /dev/null +++ b/src/prompt.rs @@ -0,0 +1,205 @@ +use std::path::Path; + +use arcstr::ArcStr; +use ndarray::{Array2, Array3, ArrayView3, Axis}; +use ort::{inputs, value::TensorRef}; +use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer, TruncationParams}; + +use crate::error::YoloError; + +#[derive(Debug, Clone)] +pub struct YoloPromptEmbeddings { + labels: Vec, + embeddings: Array3, +} + +impl YoloPromptEmbeddings { + pub fn try_new( + labels: impl IntoIterator>, + embeddings: Array3, + ) -> Result { + if embeddings.shape().len() != 3 || embeddings.shape()[0] != 1 { + return Err(YoloError::InvalidPromptEmbeddingShape(embeddings.shape().to_vec())); + } + + let labels = labels.into_iter().map(Into::into).collect::>(); + if labels.is_empty() { + return Err(YoloError::EmptyPromptList); + } + if labels.len() != embeddings.shape()[1] { + return Err(YoloError::PromptLabelEmbeddingMismatch { + labels: labels.len(), + embeddings: embeddings.shape()[1], + }); + } + + Ok(Self { labels, embeddings }) + } + + pub fn labels(&self) -> &[ArcStr] { + &self.labels + } + + pub fn embeddings(&self) -> &Array3 { + &self.embeddings + } + + pub fn view(&self) -> YoloPromptEmbeddingsView<'_> { + YoloPromptEmbeddingsView { + labels: &self.labels, + embeddings: self.embeddings.view(), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct YoloPromptEmbeddingsView<'a> { + pub labels: &'a [ArcStr], + pub embeddings: ArrayView3<'a, f32>, +} + +#[derive(Debug)] +pub struct YoloPromptEncoderSession { + session: ort::session::Session, + tokenizer: Tokenizer, + max_length: Option, +} + +impl YoloPromptEncoderSession { + pub fn from_files( + model_path: impl AsRef, + tokenizer_path: impl AsRef, + ) -> Result { + let session = ort::session::Session::builder() + .map_err(YoloError::OrtSessionBuildError)? + .commit_from_file(model_path) + .map_err(YoloError::OrtSessionLoadError)?; + + let tokenizer = Tokenizer::from_file(tokenizer_path.as_ref()) + .map_err(|error| YoloError::TokenizerLoadError(error.to_string()))?; + + Ok(Self { + session, + tokenizer, + max_length: None, + }) + } + + pub fn encode<'a>( + &mut self, + prompts: impl IntoIterator, + ) -> Result { + let labels = prompts + .into_iter() + .map(ArcStr::from) + .collect::>(); + if labels.is_empty() { + return Err(YoloError::EmptyPromptList); + } + + let prompt_texts = labels.iter().map(ToString::to_string).collect::>(); + let mut tokenizer = self.tokenizer.clone(); + if let Some(max_length) = self.max_length { + tokenizer + .with_truncation(Some(TruncationParams { + max_length, + ..Default::default() + })) + .map_err(|error| YoloError::PromptTokenizationError(error.to_string()))?; + tokenizer.with_padding(Some(PaddingParams { + strategy: PaddingStrategy::Fixed(max_length), + ..Default::default() + })); + } else { + tokenizer.with_padding(Some(PaddingParams { + strategy: PaddingStrategy::BatchLongest, + ..Default::default() + })); + } + + let encodings = tokenizer + .encode_batch(prompt_texts, true) + .map_err(|error| YoloError::PromptTokenizationError(error.to_string()))?; + let batch_size = encodings.len(); + let sequence_length = encodings.first().map(|encoding| encoding.len()).unwrap_or(0); + + let input_ids = Array2::from_shape_vec( + (batch_size, sequence_length), + encodings + .iter() + .flat_map(|encoding| encoding.get_ids().iter().map(|token| i64::from(*token))) + .collect(), + ) + .map_err(|_| YoloError::PromptTokenizationError("failed to build input_ids tensor".into()))?; + let attention_mask = Array2::from_shape_vec( + (batch_size, sequence_length), + encodings + .iter() + .flat_map(|encoding| encoding.get_attention_mask().iter().map(|token| i64::from(*token))) + .collect(), + ) + .map_err(|_| YoloError::PromptTokenizationError("failed to build attention_mask tensor".into()))?; + let token_type_ids = Array2::::zeros((batch_size, sequence_length)); + + let output_name = self + .session + .outputs() + .first() + .map(|output| output.name().to_owned()) + .ok_or(YoloError::MissingModelOutput)?; + let outputs = match self.session.inputs().len() { + 1 => { + let input_name = self.session.inputs()[0].name().to_owned(); + let inputs = inputs![input_name.as_str() => TensorRef::from_array_view(input_ids.view()).map_err(YoloError::OrtInputError)?]; + self.session.run(inputs).map_err(YoloError::OrtInferenceError)? + } + 2 => { + let input_name = self.session.inputs()[0].name().to_owned(); + let mask_name = self.session.inputs()[1].name().to_owned(); + let inputs = inputs![ + input_name.as_str() => TensorRef::from_array_view(input_ids.view()).map_err(YoloError::OrtInputError)?, + mask_name.as_str() => TensorRef::from_array_view(attention_mask.view()).map_err(YoloError::OrtInputError)?, + ]; + self.session.run(inputs).map_err(YoloError::OrtInferenceError)? + } + 3 => { + let first_name = self.session.inputs()[0].name().to_owned(); + let second_name = self.session.inputs()[1].name().to_owned(); + let third_name = self.session.inputs()[2].name().to_owned(); + let inputs = inputs![ + first_name.as_str() => TensorRef::from_array_view(input_ids.view()).map_err(YoloError::OrtInputError)?, + second_name.as_str() => TensorRef::from_array_view(attention_mask.view()).map_err(YoloError::OrtInputError)?, + third_name.as_str() => TensorRef::from_array_view(token_type_ids.view()).map_err(YoloError::OrtInputError)?, + ]; + self.session.run(inputs).map_err(YoloError::OrtInferenceError)? + } + count => return Err(YoloError::UnsupportedPromptEncoderInputCount(count)), + }; + let output = outputs[output_name.as_str()] + .try_extract_array::() + .map_err(YoloError::OrtExtractTensorError)?; + + let mut embeddings = if let Ok(output) = output.view().into_dimensionality::() { + output.to_owned().insert_axis(Axis(0)) + } else if let Ok(output) = output.view().into_dimensionality::() { + if output.shape()[0] != 1 { + return Err(YoloError::InvalidPromptEncoderOutputShape(output.shape().to_vec())); + } + output.to_owned() + } else { + return Err(YoloError::InvalidPromptEncoderOutputShape(output.shape().to_vec())); + }; + + normalize_last_axis(&mut embeddings); + YoloPromptEmbeddings::try_new(labels, embeddings) + } +} + +fn normalize_last_axis(embeddings: &mut Array3) { + for mut prompt in embeddings.axis_iter_mut(Axis(1)) { + let norm = prompt.iter().map(|value| value * value).sum::().sqrt(); + if norm > 0.0 { + prompt.iter_mut().for_each(|value| *value /= norm); + } + } +} \ No newline at end of file