[Example] YOLO-Series(v5-11) ONNXRuntime Rust (#17311)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>pull/17318/head
Jamjamjon
4 weeks ago
committed by
GitHub
parent
d28caa9a58
commit
f95dc37311
8 changed files with 362 additions and 29 deletions
@ -0,0 +1,12 @@ |
|||||||
|
[package] |
||||||
|
name = "YOLO-ONNXRuntime-Rust" |
||||||
|
version = "0.1.0" |
||||||
|
edition = "2021" |
||||||
|
authors = ["Jamjamjon <xxyydzml@outlook.com>"] |
||||||
|
|
||||||
|
[dependencies] |
||||||
|
anyhow = "1.0.92" |
||||||
|
clap = "4.5.20" |
||||||
|
tracing = "0.1.40" |
||||||
|
tracing-subscriber = "0.3.18" |
||||||
|
usls = { version = "0.0.19", features = ["auto"] } |
||||||
@ -0,0 +1,94 @@ |
|||||||
|
# YOLO-Series ONNXRuntime Rust Demo for Core YOLO Tasks |
||||||
|
|
||||||
|
This repository provides a Rust demo for key YOLO-Series tasks such as `Classification`, `Segmentation`, `Detection`, `Pose Detection`, and `OBB` using ONNXRuntime. It supports various YOLO models (v5 - 11) across multiple vision tasks. |
||||||
|
|
||||||
|
## Introduction |
||||||
|
|
||||||
|
- This example leverages the latest versions of both ONNXRuntime and YOLO models. |
||||||
|
- We utilize the [usls](https://github.com/jamjamjon/usls/tree/main) crate to streamline YOLO model inference, providing efficient data loading, visualization, and optimized inference performance. |
||||||
|
|
||||||
|
## Features |
||||||
|
|
||||||
|
- **Extensive Model Compatibility**: Supports `YOLOv5`, `YOLOv6`, `YOLOv7`, `YOLOv8`, `YOLOv9`, `YOLOv10`, `YOLO11`, `YOLO-world`, `RTDETR`, and others, covering a wide range of YOLO versions. |
||||||
|
- **Versatile Task Coverage**: Includes `Classification`, `Segmentation`, `Detection`, `Pose`, and `OBB`. |
||||||
|
- **Precision Flexibility**: Works with `FP16` and `FP32` ONNX models. |
||||||
|
- **Execution Providers**: Accelerated support for `CPU`, `CUDA`, `CoreML`, and `TensorRT`. |
||||||
|
- **Dynamic Input Shapes**: Dynamically adjusts to variable `batch`, `width`, and `height` dimensions for flexible model input. |
||||||
|
- **Flexible Data Loading**: The `DataLoader` handles images, folders, videos, and video streams. |
||||||
|
- **Real-Time Display and Video Export**: `Viewer` provides real-time frame visualization and video export functions, similar to OpenCV’s `imshow()` and `imwrite()`. |
||||||
|
- **Enhanced Annotation and Visualization**: The `Annotator` facilitates comprehensive result rendering, with support for bounding boxes (HBB), oriented bounding boxes (OBB), polygons, masks, keypoints, and text labels. |
||||||
|
|
||||||
|
## Setup Instructions |
||||||
|
|
||||||
|
### 1. ONNXRuntime Linking |
||||||
|
|
||||||
|
<details> |
||||||
|
<summary>You have two options to link the ONNXRuntime library:</summary> |
||||||
|
|
||||||
|
- **Option 1: Manual Linking** |
||||||
|
|
||||||
|
- For detailed setup, consult the [ONNX Runtime linking documentation](https://ort.pyke.io/setup/linking). |
||||||
|
- **Linux or macOS**: |
||||||
|
1. Download the ONNX Runtime package from the [Releases page](https://github.com/microsoft/onnxruntime/releases). |
||||||
|
2. Set up the library path by exporting the `ORT_DYLIB_PATH` environment variable: |
||||||
|
```shell |
||||||
|
export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.19.0 |
||||||
|
``` |
||||||
|
|
||||||
|
- **Option 2: Automatic Download** |
||||||
|
- Use the `--features auto` flag to handle downloading automatically: |
||||||
|
```shell |
||||||
|
cargo run -r --example yolo --features auto |
||||||
|
``` |
||||||
|
|
||||||
|
</details> |
||||||
|
|
||||||
|
### 2. \[Optional\] Install CUDA, CuDNN, and TensorRT |
||||||
|
|
||||||
|
- The CUDA execution provider requires CUDA version `12.x`. |
||||||
|
- The TensorRT execution provider requires both CUDA `12.x` and TensorRT `10.x`. |
||||||
|
|
||||||
|
### 3. \[Optional\] Install ffmpeg |
||||||
|
|
||||||
|
To view video frames and save video inferences, install `rust-ffmpeg`. For instructions, see: |
||||||
|
[https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building#dependencies](https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building#dependencies) |
||||||
|
|
||||||
|
## Get Started |
||||||
|
|
||||||
|
```Shell |
||||||
|
# customized |
||||||
|
cargo run -r -- --task detect --ver v8 --nc 6 --model xxx.onnx # YOLOv8 |
||||||
|
|
||||||
|
# Classify |
||||||
|
cargo run -r -- --task classify --ver v5 --scale s --width 224 --height 224 --nc 1000 # YOLOv5 |
||||||
|
cargo run -r -- --task classify --ver v8 --scale n --width 224 --height 224 --nc 1000 # YOLOv8 |
||||||
|
cargo run -r -- --task classify --ver v11 --scale n --width 224 --height 224 --nc 1000 # YOLOv11 |
||||||
|
|
||||||
|
# Detect |
||||||
|
cargo run -r -- --task detect --ver v5 --scale n # YOLOv5 |
||||||
|
cargo run -r -- --task detect --ver v6 --scale n # YOLOv6 |
||||||
|
cargo run -r -- --task detect --ver v7 --scale t # YOLOv7 |
||||||
|
cargo run -r -- --task detect --ver v8 --scale n # YOLOv8 |
||||||
|
cargo run -r -- --task detect --ver v9 --scale t # YOLOv9 |
||||||
|
cargo run -r -- --task detect --ver v10 --scale n # YOLOv10 |
||||||
|
cargo run -r -- --task detect --ver v11 --scale n # YOLOv11 |
||||||
|
cargo run -r -- --task detect --ver rtdetr --scale l # RTDETR |
||||||
|
|
||||||
|
# Pose |
||||||
|
cargo run -r -- --task pose --ver v8 --scale n # YOLOv8-Pose |
||||||
|
cargo run -r -- --task pose --ver v11 --scale n # YOLOv11-Pose |
||||||
|
|
||||||
|
# Segment |
||||||
|
cargo run -r -- --task segment --ver v5 --scale n # YOLOv5-Segment |
||||||
|
cargo run -r -- --task segment --ver v8 --scale n # YOLOv8-Segment |
||||||
|
cargo run -r -- --task segment --ver v11 --scale n # YOLOv8-Segment |
||||||
|
cargo run -r -- --task segment --ver v8 --model yolo/FastSAM-s-dyn-f16.onnx # FastSAM |
||||||
|
|
||||||
|
# OBB |
||||||
|
cargo run -r -- --ver v8 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv8-Obb |
||||||
|
cargo run -r -- --ver v11 --task obb --scale n --width 1024 --height 1024 --source images/dota.png # YOLOv11-Obb |
||||||
|
``` |
||||||
|
|
||||||
|
**`cargo run -- --help` for more options** |
||||||
|
|
||||||
|
For more details, please refer to [usls-yolo](https://github.com/jamjamjon/usls/tree/main/examples/yolo). |
||||||
@ -0,0 +1,236 @@ |
|||||||
|
use anyhow::Result; |
||||||
|
use clap::Parser; |
||||||
|
|
||||||
|
use usls::{ |
||||||
|
models::YOLO, Annotator, DataLoader, Device, Options, Viewer, Vision, YOLOScale, YOLOTask, |
||||||
|
YOLOVersion, COCO_SKELETONS_16, |
||||||
|
}; |
||||||
|
|
||||||
|
#[derive(Parser, Clone)] |
||||||
|
#[command(author, version, about, long_about = None)] |
||||||
|
pub struct Args { |
||||||
|
/// Path to the ONNX model
|
||||||
|
#[arg(long)] |
||||||
|
pub model: Option<String>, |
||||||
|
|
||||||
|
/// Input source path
|
||||||
|
#[arg(long, default_value_t = String::from("../../ultralytics/assets/bus.jpg"))] |
||||||
|
pub source: String, |
||||||
|
|
||||||
|
/// YOLO Task
|
||||||
|
#[arg(long, value_enum, default_value_t = YOLOTask::Detect)] |
||||||
|
pub task: YOLOTask, |
||||||
|
|
||||||
|
/// YOLO Version
|
||||||
|
#[arg(long, value_enum, default_value_t = YOLOVersion::V8)] |
||||||
|
pub ver: YOLOVersion, |
||||||
|
|
||||||
|
/// YOLO Scale
|
||||||
|
#[arg(long, value_enum, default_value_t = YOLOScale::N)] |
||||||
|
pub scale: YOLOScale, |
||||||
|
|
||||||
|
/// Batch size
|
||||||
|
#[arg(long, default_value_t = 1)] |
||||||
|
pub batch_size: usize, |
||||||
|
|
||||||
|
/// Minimum input width
|
||||||
|
#[arg(long, default_value_t = 224)] |
||||||
|
pub width_min: isize, |
||||||
|
|
||||||
|
/// Input width
|
||||||
|
#[arg(long, default_value_t = 640)] |
||||||
|
pub width: isize, |
||||||
|
|
||||||
|
/// Maximum input width
|
||||||
|
#[arg(long, default_value_t = 1024)] |
||||||
|
pub width_max: isize, |
||||||
|
|
||||||
|
/// Minimum input height
|
||||||
|
#[arg(long, default_value_t = 224)] |
||||||
|
pub height_min: isize, |
||||||
|
|
||||||
|
/// Input height
|
||||||
|
#[arg(long, default_value_t = 640)] |
||||||
|
pub height: isize, |
||||||
|
|
||||||
|
/// Maximum input height
|
||||||
|
#[arg(long, default_value_t = 1024)] |
||||||
|
pub height_max: isize, |
||||||
|
|
||||||
|
/// Number of classes
|
||||||
|
#[arg(long, default_value_t = 80)] |
||||||
|
pub nc: usize, |
||||||
|
|
||||||
|
/// Class confidence
|
||||||
|
#[arg(long)] |
||||||
|
pub confs: Vec<f32>, |
||||||
|
|
||||||
|
/// Enable TensorRT support
|
||||||
|
#[arg(long)] |
||||||
|
pub trt: bool, |
||||||
|
|
||||||
|
/// Enable CUDA support
|
||||||
|
#[arg(long)] |
||||||
|
pub cuda: bool, |
||||||
|
|
||||||
|
/// Enable CoreML support
|
||||||
|
#[arg(long)] |
||||||
|
pub coreml: bool, |
||||||
|
|
||||||
|
/// Use TensorRT half precision
|
||||||
|
#[arg(long)] |
||||||
|
pub half: bool, |
||||||
|
|
||||||
|
/// Device ID to use
|
||||||
|
#[arg(long, default_value_t = 0)] |
||||||
|
pub device_id: usize, |
||||||
|
|
||||||
|
/// Enable performance profiling
|
||||||
|
#[arg(long)] |
||||||
|
pub profile: bool, |
||||||
|
|
||||||
|
/// Disable contour drawing, for saving time
|
||||||
|
#[arg(long)] |
||||||
|
pub no_contours: bool, |
||||||
|
|
||||||
|
/// Show result
|
||||||
|
#[arg(long)] |
||||||
|
pub view: bool, |
||||||
|
|
||||||
|
/// Do not save output
|
||||||
|
#[arg(long)] |
||||||
|
pub nosave: bool, |
||||||
|
} |
||||||
|
|
||||||
|
fn main() -> Result<()> { |
||||||
|
let args = Args::parse(); |
||||||
|
|
||||||
|
// logger
|
||||||
|
if args.profile { |
||||||
|
tracing_subscriber::fmt() |
||||||
|
.with_max_level(tracing::Level::INFO) |
||||||
|
.init(); |
||||||
|
} |
||||||
|
|
||||||
|
// model path
|
||||||
|
let path = match &args.model { |
||||||
|
None => format!( |
||||||
|
"yolo/{}-{}-{}.onnx", |
||||||
|
args.ver.name(), |
||||||
|
args.scale.name(), |
||||||
|
args.task.name() |
||||||
|
), |
||||||
|
Some(x) => x.to_string(), |
||||||
|
}; |
||||||
|
|
||||||
|
// saveout
|
||||||
|
let saveout = match &args.model { |
||||||
|
None => format!( |
||||||
|
"{}-{}-{}", |
||||||
|
args.ver.name(), |
||||||
|
args.scale.name(), |
||||||
|
args.task.name() |
||||||
|
), |
||||||
|
Some(x) => { |
||||||
|
let p = std::path::PathBuf::from(&x); |
||||||
|
p.file_stem().unwrap().to_str().unwrap().to_string() |
||||||
|
} |
||||||
|
}; |
||||||
|
|
||||||
|
// device
|
||||||
|
let device = if args.cuda { |
||||||
|
Device::Cuda(args.device_id) |
||||||
|
} else if args.trt { |
||||||
|
Device::Trt(args.device_id) |
||||||
|
} else if args.coreml { |
||||||
|
Device::CoreML(args.device_id) |
||||||
|
} else { |
||||||
|
Device::Cpu(args.device_id) |
||||||
|
}; |
||||||
|
|
||||||
|
// build options
|
||||||
|
let options = Options::new() |
||||||
|
.with_model(&path)? |
||||||
|
.with_yolo_version(args.ver) |
||||||
|
.with_yolo_task(args.task) |
||||||
|
.with_device(device) |
||||||
|
.with_trt_fp16(args.half) |
||||||
|
.with_ixx(0, 0, (1, args.batch_size as _, 4).into()) |
||||||
|
.with_ixx(0, 2, (args.height_min, args.height, args.height_max).into()) |
||||||
|
.with_ixx(0, 3, (args.width_min, args.width, args.width_max).into()) |
||||||
|
.with_confs(if args.confs.is_empty() { |
||||||
|
&[0.2, 0.15] |
||||||
|
} else { |
||||||
|
&args.confs |
||||||
|
}) |
||||||
|
.with_nc(args.nc) |
||||||
|
.with_find_contours(!args.no_contours) // find contours or not
|
||||||
|
// .with_names(&COCO_CLASS_NAMES_80) // detection class names
|
||||||
|
// .with_names2(&COCO_KEYPOINTS_17) // keypoints class names
|
||||||
|
// .exclude_classes(&[0])
|
||||||
|
// .retain_classes(&[0, 5])
|
||||||
|
.with_profile(args.profile); |
||||||
|
|
||||||
|
// build model
|
||||||
|
let mut model = YOLO::new(options)?; |
||||||
|
|
||||||
|
// build dataloader
|
||||||
|
let dl = DataLoader::new(&args.source)? |
||||||
|
.with_batch(model.batch() as _) |
||||||
|
.build()?; |
||||||
|
|
||||||
|
// build annotator
|
||||||
|
let annotator = Annotator::default() |
||||||
|
.with_skeletons(&COCO_SKELETONS_16) |
||||||
|
.without_masks(true) // no masks plotting when doing segment task
|
||||||
|
.with_bboxes_thickness(3) |
||||||
|
.with_keypoints_name(false) // enable keypoints names
|
||||||
|
.with_saveout_subs(&["YOLO"]) |
||||||
|
.with_saveout(&saveout); |
||||||
|
|
||||||
|
// build viewer
|
||||||
|
let mut viewer = if args.view { |
||||||
|
Some(Viewer::new().with_delay(5).with_scale(1.).resizable(true)) |
||||||
|
} else { |
||||||
|
None |
||||||
|
}; |
||||||
|
|
||||||
|
// run & annotate
|
||||||
|
for (xs, _paths) in dl { |
||||||
|
let ys = model.forward(&xs, args.profile)?; |
||||||
|
let images_plotted = annotator.plot(&xs, &ys, !args.nosave)?; |
||||||
|
|
||||||
|
// show image
|
||||||
|
match &mut viewer { |
||||||
|
Some(viewer) => viewer.imshow(&images_plotted)?, |
||||||
|
None => continue, |
||||||
|
} |
||||||
|
|
||||||
|
// check out window and key event
|
||||||
|
match &mut viewer { |
||||||
|
Some(viewer) => { |
||||||
|
if !viewer.is_open() || viewer.is_key_pressed(usls::Key::Escape) { |
||||||
|
break; |
||||||
|
} |
||||||
|
} |
||||||
|
None => continue, |
||||||
|
} |
||||||
|
|
||||||
|
// write video
|
||||||
|
if !args.nosave { |
||||||
|
match &mut viewer { |
||||||
|
Some(viewer) => viewer.write_batch(&images_plotted)?, |
||||||
|
None => continue, |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
// finish video write
|
||||||
|
if !args.nosave { |
||||||
|
if let Some(viewer) = &mut viewer { |
||||||
|
viewer.finish_write()?; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
Ok(()) |
||||||
|
} |
||||||
Loading…
Reference in new issue