[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
3 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