feat: support yolov10 (#15)

Signed-off-by: oatiz <[email protected]>

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) |      ✅      |      ✅      |             ✅             |            ✅            |

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)

examples/yolov10/main.rs

@@ -0,0 +1,28 @@
+use usls::{
+    models::{YOLOVersion, YOLO},
+    Annotator, DataLoader, Options,
+};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // 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(())
+}

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<f32>,
     pub unclip_ratio: f32, // DB
     pub yolo_task: Option<YOLOTask>,
+    pub yolo_version: Option<YOLOVersion>,
     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

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;

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<Array<f32, IxDyn>>, xs0: &[DynamicImage]) -> Result<Vec<Y>> {
@@ -411,6 +425,59 @@ impl YOLO {
         Ok(ys)
     }
 
+    pub fn postprocess_v10(
+        &self,
+        xs: Vec<Array<f32, IxDyn>>,
+        xs0: &[DynamicImage],
+    ) -> Result<Vec<Y>> {
+        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
     }