feat: initial HSAP platform
Huaxu Sentinel Active Safety Platform with embedded algorithm code, Docker Compose setup, and vendored dataset scaffolds for clone-and-run. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,14 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
[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,100 @@
|
||||
# YOLO-Series ONNXRuntime Rust Demo for Core YOLO Tasks
|
||||
|
||||
This repository provides a [Rust](https://rust-lang.org/) demo showcasing key [Ultralytics YOLO](https://docs.ultralytics.com/) series tasks such as [Classification](https://docs.ultralytics.com/tasks/classify/), [Segmentation](https://docs.ultralytics.com/tasks/segment/), [Detection](https://docs.ultralytics.com/tasks/detect/), [Pose Estimation](https://docs.ultralytics.com/tasks/pose/), and Oriented Bounding Box ([OBB](https://docs.ultralytics.com/tasks/obb/)) detection using the [ONNXRuntime](https://github.com/microsoft/onnxruntime). It supports various YOLO models (YOLOv5 through YOLO11) across multiple computer vision tasks.
|
||||
|
||||
## ✨ Introduction
|
||||
|
||||
- This example leverages the latest versions of both the [ONNX Runtime](https://onnxruntime.ai/) and popular YOLO models.
|
||||
- We utilize the [usls crate](https://github.com/jamjamjon/usls/tree/main) to streamline YOLO model inference in Rust, providing efficient data loading, visualization, and optimized inference performance. This allows developers to easily integrate state-of-the-art object detection into their Rust applications.
|
||||
|
||||
## 🚀 Features
|
||||
|
||||
- **Extensive Model Compatibility**: Supports a wide range of YOLO versions including [YOLOv5](https://docs.ultralytics.com/models/yolov5/), [YOLOv6](https://docs.ultralytics.com/models/yolov6/), [YOLOv7](https://docs.ultralytics.com/models/yolov7/), [YOLOv8](https://docs.ultralytics.com/models/yolov8/), [YOLOv9](https://docs.ultralytics.com/models/yolov9/), [YOLOv10](https://docs.ultralytics.com/models/yolov10/), [YOLO11](https://docs.ultralytics.com/models/yolo11/), [YOLO-World](https://docs.ultralytics.com/models/yolo-world/), [RT-DETR](https://docs.ultralytics.com/models/rtdetr/), and others.
|
||||
- **Versatile Task Coverage**: Includes examples for `Classification`, `Segmentation`, `Detection`, `Pose`, and `OBB`.
|
||||
- **Precision Flexibility**: Works seamlessly with `FP16` and `FP32` precision [ONNX models](https://docs.ultralytics.com/integrations/onnx/).
|
||||
- **Execution Providers**: Accelerated support for `CPU`, [CUDA](https://developer.nvidia.com/cuda-toolkit), [CoreML](https://developer.apple.com/documentation/coreml), and [TensorRT](https://docs.ultralytics.com/integrations/tensorrt/).
|
||||
- **Dynamic Input Shapes**: Dynamically adjusts to variable `batch`, `width`, and `height` dimensions for flexible model input.
|
||||
- **Flexible Data Loading**: The `DataLoader` component handles images, folders, videos, and real-time video streams.
|
||||
- **Real-Time Display and Video Export**: The `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, supporting 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 instructions, consult the [ONNX Runtime linking documentation](https://ort.pyke.io/setup/linking).
|
||||
- **Linux or macOS**:
|
||||
1. Download the appropriate ONNX Runtime package from the official [Releases page](https://github.com/microsoft/onnxruntime/releases).
|
||||
2. Set up the library path by exporting the `ORT_DYLIB_PATH` environment variable, pointing to the downloaded library file:
|
||||
```bash
|
||||
# Example path, replace with your actual path
|
||||
export ORT_DYLIB_PATH=/path/to/onnxruntime/lib/libonnxruntime.so.1.19.0
|
||||
```
|
||||
|
||||
- **Option 2: Automatic Download**
|
||||
- Use the `--features auto` flag with Cargo to let the build script handle downloading the library automatically:
|
||||
```bash
|
||||
cargo run -r --example yolo --features auto
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### 2. [Optional] Install CUDA, CuDNN, and TensorRT
|
||||
|
||||
- The CUDA execution provider requires [NVIDIA CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit) version `12.x`.
|
||||
- The TensorRT execution provider requires both CUDA `12.x` and [NVIDIA TensorRT](https://developer.nvidia.com/tensorrt) `10.x`. Ensure [cuDNN](https://developer.nvidia.com/cudnn) is also correctly installed.
|
||||
|
||||
### 3. [Optional] Install ffmpeg
|
||||
|
||||
To enable viewing video frames and saving video inferences, install the `rust-ffmpeg` crate's dependencies. Follow the instructions provided here:
|
||||
[https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building#dependencies](https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building#dependencies)
|
||||
|
||||
## ▶️ Get Started
|
||||
|
||||
Run the examples using Cargo. The `--` separates Cargo arguments from the example's arguments.
|
||||
|
||||
```bash
|
||||
# Run a custom model (e.g., YOLOv8 detection)
|
||||
cargo run -r -- --task detect --ver v8 --nc 6 --model path/to/your/model.onnx
|
||||
|
||||
# Classify examples
|
||||
cargo run -r -- --task classify --ver v5 --scale s --width 224 --height 224 --nc 1000 # YOLOv5 Classification
|
||||
cargo run -r -- --task classify --ver v8 --scale n --width 224 --height 224 --nc 1000 # YOLOv8 Classification
|
||||
cargo run -r -- --task classify --ver v11 --scale n --width 224 --height 224 --nc 1000 # YOLO11 Classification
|
||||
|
||||
# Detect examples
|
||||
cargo run -r -- --task detect --ver v5 --scale n # YOLOv5 Detection
|
||||
cargo run -r -- --task detect --ver v6 --scale n # YOLOv6 Detection
|
||||
cargo run -r -- --task detect --ver v7 --scale t # YOLOv7 Detection
|
||||
cargo run -r -- --task detect --ver v8 --scale n # YOLOv8 Detection
|
||||
cargo run -r -- --task detect --ver v9 --scale t # YOLOv9 Detection
|
||||
cargo run -r -- --task detect --ver v10 --scale n # YOLOv10 Detection
|
||||
cargo run -r -- --task detect --ver v11 --scale n # YOLO11 Detection
|
||||
cargo run -r -- --task detect --ver rtdetr --scale l # RT-DETR Detection
|
||||
|
||||
# Pose examples
|
||||
cargo run -r -- --task pose --ver v8 --scale n # YOLOv8-Pose Estimation
|
||||
cargo run -r -- --task pose --ver v11 --scale n # YOLO11-Pose Estimation
|
||||
|
||||
# Segment examples
|
||||
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 # YOLO11-Segment
|
||||
cargo run -r -- --task segment --ver v8 --model path/to/FastSAM-s-dyn-f16.onnx # FastSAM Segmentation
|
||||
|
||||
# OBB (Oriented Bounding Box) examples
|
||||
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 # YOLO11-OBB
|
||||
```
|
||||
|
||||
**Use `cargo run -- --help` to see all available options.**
|
||||
|
||||
For more detailed information and advanced usage, please refer to the [usls-yolo example documentation](https://github.com/jamjamjon/usls/tree/main/examples/yolo).
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! If you'd like to improve this demo or add new features, please feel free to submit issues or pull requests on the repository. Your input helps make the Ultralytics ecosystem better for everyone. Check out the [Ultralytics Contribution Guide](https://docs.ultralytics.com/help/contributing/) for more details.
|
||||
@@ -0,0 +1,238 @@
|
||||
// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user