diff --git a/README.md b/README.md index 37be300..ab8cf78 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # usls -A Rust library integrated with **ONNXRuntime**, providing a collection of **Computer Vison** and **Vision-Language** models including [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [RTDETR](https://arxiv.org/abs/2304.08069), [CLIP](https://github.com/openai/CLIP), [DINOv2](https://github.com/facebookresearch/dinov2), [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [BLIP](https://arxiv.org/abs/2201.12086), [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR), [Depth-Anything](https://github.com/LiheYoung/Depth-Anything), [MODNet](https://github.com/ZHKKKe/MODNet) and others. +A Rust library integrated with **ONNXRuntime**, providing a collection of **Computer Vison** and **Vision-Language** models including [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv8](https://github.com/ultralytics/ultralytics), [YOLOv9](https://github.com/WongKinYiu/yolov9), [YOLOv10](https://github.com/THU-MIG/yolov10) [RTDETR](https://arxiv.org/abs/2304.08069), [CLIP](https://github.com/openai/CLIP), [DINOv2](https://github.com/facebookresearch/dinov2), [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM), [YOLO-World](https://github.com/AILab-CVC/YOLO-World), [BLIP](https://arxiv.org/abs/2201.12086), [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR), [Depth-Anything](https://github.com/LiheYoung/Depth-Anything), [MODNet](https://github.com/ZHKKKe/MODNet) and others. ## Recently Updated @@ -37,6 +37,7 @@ A Rust library integrated with **ONNXRuntime**, providing a collection of **Comp | [YOLOv8-classification](https://github.com/ultralytics/ultralytics) | Classification | [demo](examples/yolov8) | ✅ | ✅ | ✅ | ✅ | | [YOLOv8-segmentation](https://github.com/ultralytics/ultralytics) | Instance Segmentation | [demo](examples/yolov8) | ✅ | ✅ | ✅ | ✅ | | [YOLOv9](https://github.com/WongKinYiu/yolov9) | Object Detection | [demo](examples/yolov9) | ✅ | ✅ | ✅ | ✅ | +| [YOLOv10](https://github.com/THU-MIG/yolov10) | Object Detection | [demo](examples/yolov10) | ✅ | ✅ | ✅ | ✅ | | [RT-DETR](https://arxiv.org/abs/2304.08069) | Object Detection | [demo](examples/rtdetr) | ✅ | ✅ | ✅ | ✅ | | [FastSAM](https://github.com/CASIA-IVA-Lab/FastSAM) | Instance Segmentation | [demo](examples/fastsam) | ✅ | ✅ | ✅ | ✅ | | [YOLO-World](https://github.com/AILab-CVC/YOLO-World) | Object Detection | [demo](examples/yolo-world) | ✅ | ✅ | ✅ | ✅ | diff --git a/examples/yolov10/README.md b/examples/yolov10/README.md new file mode 100644 index 0000000..dab69e4 --- /dev/null +++ b/examples/yolov10/README.md @@ -0,0 +1,26 @@ +## Quick Start + +```shell +cargo run -r --example yolov10 +``` + +## Export ONNX Model + +- **Export** + + ```shell + # clone repo and install dependencies + git clone https://github.com/THU-MIG/yolov10.git + cd yolov10 + pip install -r requirements.txt + + # donwload `pt` weights + wget https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10n.pt + + # export ONNX model + yolo export model=yolov10n.pt format=onnx opset=13 simplify dynamic + ``` + +## Results + +![](./demo.png) diff --git a/examples/yolov10/demo.png b/examples/yolov10/demo.png new file mode 100644 index 0000000..0cd261c Binary files /dev/null and b/examples/yolov10/demo.png differ diff --git a/examples/yolov10/main.rs b/examples/yolov10/main.rs new file mode 100644 index 0000000..8ce8a5e --- /dev/null +++ b/examples/yolov10/main.rs @@ -0,0 +1,28 @@ +use usls::{ + models::{YOLOVersion, YOLO}, + Annotator, DataLoader, Options, +}; + +fn main() -> Result<(), Box> { + // build model + let options = Options::default() + .with_model("yolov10n-dyn.onnx")? + .with_yolo_version(YOLOVersion::V10) + .with_i00((1, 1, 4).into()) + .with_i02((416, 640, 800).into()) + .with_i03((416, 640, 800).into()) + .with_confs(&[0.4, 0.15]); + let mut model = YOLO::new(options)?; + + // load image + let x = vec![DataLoader::try_read("./assets/bus.jpg")?]; + + // run + let y = model.run(&x)?; + + // annotate + let annotator = Annotator::default().with_saveout("YOLOv10"); + annotator.annotate(&x, &y); + + Ok(()) +} diff --git a/src/core/options.rs b/src/core/options.rs index 9c350b6..9faa92e 100644 --- a/src/core/options.rs +++ b/src/core/options.rs @@ -1,6 +1,10 @@ use anyhow::Result; -use crate::{auto_load, models::YOLOTask, Device, MinOptMax}; +use crate::{ + auto_load, + models::{YOLOTask, YOLOVersion}, + Device, MinOptMax, +}; /// Options for building models #[derive(Debug, Clone)] @@ -56,6 +60,7 @@ pub struct Options { pub min_height: Option, pub unclip_ratio: f32, // DB pub yolo_task: Option, + pub yolo_version: Option, pub anchors_first: bool, // yolo model output format like: [batch_size, anchors, xywh_clss_xxx] pub conf_independent: bool, // xywh_conf_clss pub apply_probs_softmax: bool, @@ -111,6 +116,7 @@ impl Default for Options { min_height: None, unclip_ratio: 1.5, yolo_task: None, + yolo_version: None, anchors_first: false, conf_independent: false, apply_probs_softmax: false, @@ -159,6 +165,11 @@ impl Options { self } + pub fn with_yolo_version(mut self, x: YOLOVersion) -> Self { + self.yolo_version = Some(x); + self + } + pub fn with_conf_independent(mut self, x: bool) -> Self { self.conf_independent = x; self diff --git a/src/models/mod.rs b/src/models/mod.rs index 5b521f6..2815614 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -19,5 +19,5 @@ pub use modnet::MODNet; pub use rtdetr::RTDETR; pub use rtmo::RTMO; pub use svtr::SVTR; -pub use yolo::{YOLOTask, YOLO}; +pub use yolo::{YOLOTask, YOLOVersion, YOLO}; pub use yolop::YOLOPv2; diff --git a/src/models/yolo.rs b/src/models/yolo.rs index 6e56309..9721ddb 100644 --- a/src/models/yolo.rs +++ b/src/models/yolo.rs @@ -18,6 +18,12 @@ pub enum YOLOTask { Obb, } +#[derive(Debug, Copy, Clone, ValueEnum)] +pub enum YOLOVersion { + V8, + V10, +} + #[derive(Debug)] pub struct YOLO { engine: OrtEngine, @@ -28,6 +34,7 @@ pub struct YOLO { width: MinOptMax, batch: MinOptMax, task: YOLOTask, + version: YOLOVersion, confs: DynConf, kconfs: DynConf, iou: f32, @@ -64,6 +71,8 @@ impl YOLO { }, }; + let version = options.yolo_version.unwrap_or(YOLOVersion::V8); + // try from custom class names, and then model metadata let mut names = options.names.or(Self::fetch_names(&engine)); let nc = match options.nc { @@ -121,6 +130,7 @@ impl YOLO { width, batch, task, + version, names, names_kpt, anchors_first: options.anchors_first, @@ -144,7 +154,11 @@ impl YOLO { }; let xs_ = ops::normalize(xs_, 0., 255.); let ys = self.engine.run(&[xs_])?; - self.postprocess(ys, xs) + + match self.version { + YOLOVersion::V10 => self.postprocess_v10(ys, xs), + _ => self.postprocess(ys, xs), + } } pub fn postprocess(&self, xs: Vec>, xs0: &[DynamicImage]) -> Result> { @@ -411,6 +425,59 @@ impl YOLO { Ok(ys) } + pub fn postprocess_v10( + &self, + xs: Vec>, + xs0: &[DynamicImage], + ) -> Result> { + let mut ys = Vec::new(); + for (idx, preds) in xs[0].axis_iter(Axis(0)).enumerate() { + let image_width = xs0[idx].width() as f32; + let image_height = xs0[idx].height() as f32; + match self.task { + YOLOTask::Detect => { + let ratio = (self.width() as f32 / image_width) + .min(self.height() as f32 / image_height); + + let mut y_bboxes = vec![]; + for (i, pred) in preds.axis_iter(Axis(0)).enumerate() { + let confidence = pred[CXYWH_OFFSET]; + if confidence < self.confs[0] { + continue; + } + let class_id = pred[CXYWH_OFFSET + 1] as isize; + + let bbox = pred.slice(s![0..CXYWH_OFFSET]); + // re-scale + let x = bbox[0] / ratio; + let y = bbox[1] / ratio; + let x2 = bbox[2] / ratio; + let y2 = bbox[3] / ratio; + let w = x2 - x; + let h = y2 - y; + + let y_bbox = Bbox::default() + .with_xywh(x, y, w, h) + .with_confidence(confidence) + .with_id(class_id) + .with_id_born(i as isize) + .with_name( + self.names + .as_ref() + .map(|names| names[class_id as usize].to_owned()), + ); + y_bboxes.push(y_bbox); + } + let y = Y::default().with_bboxes(&y_bboxes); + ys.push(y); + } + _ => todo!("YOLO_V10 Not supported: {:?}", self.task), + } + } + + Ok(ys) + } + pub fn batch(&self) -> isize { self.batch.opt }