update

Cargo.toml

@@ -44,4 +44,15 @@ ab_glyph = "0.2.23"
 geo = "0.28.0"
 prost = "0.12.4"
 human_bytes = "0.4.3"
-fast_image_resize = { version = "4.2.0", features = ["image"]}
+fast_image_resize = { version = "4.2.1", features = ["image"]}
+
+
+[dev-dependencies]
+criterion = "0.5.1"
+
+[[bench]]
+name = "yolo"
+harness = false
+
+[lib]
+bench = false

README.md

@@ -1,6 +1,9 @@
 # usls
 
-[![Static Badge](https://img.shields.io/crates/v/usls.svg?style=for-the-badge&logo=rust)](https://crates.io/crates/usls) ![Static Badge](https://img.shields.io/crates/d/usls?style=for-the-badge) [![Static Badge](https://img.shields.io/badge/Documents-usls-blue?style=for-the-badge&logo=docs.rs)](https://docs.rs/usls) [![Static Badge](https://img.shields.io/badge/GitHub-black?style=for-the-badge&logo=github)](https://github.com/jamjamjon/usls)
+[![Static Badge](https://img.shields.io/crates/v/usls.svg?style=for-the-badge&logo=rust)](https://crates.io/crates/usls) [![Static Badge](https://img.shields.io/badge/ONNXRuntime-v1.17.x-yellow?style=for-the-badge&logo=docs.rs)](https://github.com/microsoft/onnxruntime/releases) [![Static Badge](https://img.shields.io/badge/CUDA-11.x-green?style=for-the-badge&logo=docs.rs)](https://developer.nvidia.com/cuda-toolkit-archive) [![Static Badge](https://img.shields.io/badge/TRT-8.6.x.x-blue?style=for-the-badge&logo=docs.rs)](https://developer.nvidia.com/tensorrt)  
+[![Static Badge](https://img.shields.io/badge/Documents-usls-blue?style=for-the-badge&logo=docs.rs)](https://docs.rs/usls) ![Static Badge](https://img.shields.io/crates/d/usls?style=for-the-badge)
+
+
 
 A Rust library integrated with **ONNXRuntime**, providing a collection of **Computer Vison** and **Vision-Language** models including [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [YOLOv7](https://github.com/WongKinYiu/yolov7), [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.
 

benches/yolo.rs

@@ -0,0 +1,105 @@
+use anyhow::Result;
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+use usls::{coco, models::YOLO, DataLoader, Options, Vision, YOLOTask, YOLOVersion};
+
+// fn yolo_bench(model: &mut YOLO, xs: &[image::DynamicImage]) {
+//     let _ys = model.run(&xs).unwrap();
+// }
+
+enum Stage {
+    Pre,
+    Run,
+    Post,
+    Pipeline,
+}
+
+fn yolo_stage_bench(
+    model: &mut YOLO,
+    x: &[image::DynamicImage],
+    stage: Stage,
+    n: u64,
+) -> std::time::Duration {
+    let mut t_pre = std::time::Duration::new(0, 0);
+    let mut t_run = std::time::Duration::new(0, 0);
+    let mut t_post = std::time::Duration::new(0, 0);
+    let mut t_pipeline = std::time::Duration::new(0, 0);
+    for _ in 0..n {
+        let t0 = std::time::Instant::now();
+        let xs = model.preprocess(&x).unwrap();
+        t_pre += t0.elapsed();
+
+        let t = std::time::Instant::now();
+        let xs = model.inference(xs).unwrap();
+        t_run += t.elapsed();
+
+        let t = std::time::Instant::now();
+        let _ys = black_box(model.postprocess(xs, &x).unwrap());
+        let t1 = t.elapsed();
+        t_post += t1;
+        t_pipeline = t1;
+    }
+    match stage {
+        Stage::Pre => t_pre,
+        Stage::Run => t_run,
+        Stage::Post => t_post,
+        Stage::Pipeline => t_pipeline,
+    }
+}
+
+pub fn benchmark_cuda(c: &mut Criterion, h: isize, w: isize) -> Result<()> {
+    let mut group = c.benchmark_group(format!("YOLO ({}-{})", w, h));
+    group
+        .significance_level(0.05)
+        .sample_size(80)
+        .measurement_time(std::time::Duration::new(20, 0));
+
+    let options = Options::default()
+        .with_yolo_version(YOLOVersion::V8) // YOLOVersion: V5, V6, V7, V8, V9, V10, RTDETR
+        .with_yolo_task(YOLOTask::Detect) // YOLOTask: Classify, Detect, Pose, Segment, Obb
+        .with_model("yolov8m-dyn.onnx")?
+        .with_cuda(0)
+        // .with_cpu()
+        .with_dry_run(0)
+        .with_i00((1, 1, 4).into())
+        .with_i02((320, h, 1280).into())
+        .with_i03((320, w, 1280).into())
+        .with_confs(&[0.2, 0.15]) // class_0: 0.4, others: 0.15
+        .with_names2(&coco::KEYPOINTS_NAMES_17);
+    let mut model = YOLO::new(options)?;
+
+    let xs = vec![DataLoader::try_read("./assets/bus.jpg")?];
+
+    // group.bench_function("whole-pipeline", |b| {
+    //     b.iter(|| yolo_bench(black_box(&mut model), black_box(&xs)))
+    // });
+
+    group.bench_function("pre-process", |b| {
+        b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Pre, n))
+    });
+
+    group.bench_function("run", |b| {
+        b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Run, n))
+    });
+
+    group.bench_function("post-process", |b| {
+        b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Post, n))
+    });
+
+    group.bench_function("pipeline", |b| {
+        b.iter_custom(|n| yolo_stage_bench(&mut model, &xs, Stage::Pipeline, n))
+    });
+
+    group.finish();
+    Ok(())
+}
+
+pub fn criterion_benchmark(c: &mut Criterion) {
+    // benchmark_cuda(c, 416, 416).unwrap();
+    benchmark_cuda(c, 640, 640).unwrap();
+    benchmark_cuda(c, 448, 768).unwrap();
+    // benchmark_cuda(c, 800, 800).unwrap();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

examples/yolo/main.rs

@@ -163,7 +163,7 @@ fn main() -> Result<()> {
     // build annotator
     let annotator = Annotator::default()
         .with_skeletons(&coco::SKELETONS_16)
-        .with_bboxes_thickness(7)
+        .with_bboxes_thickness(4)
         .without_masks(true) // No masks plotting.
         .with_saveout("YOLO-Series");
 

src/core/dataloader.rs

@@ -53,42 +53,31 @@ impl Default for DataLoader {
 }
 
 impl DataLoader {
-    pub fn load<P: AsRef<Path>>(&mut self, source: P) -> Result<Self> {
-        let source = source.as_ref();
-        let mut paths = VecDeque::new();
-
-        match source {
-            s if s.is_file() => paths.push_back(s.to_path_buf()),
-            s if s.is_dir() => {
-                for entry in WalkDir::new(s)
-                    .into_iter()
-                    .filter_entry(|e| !Self::_is_hidden(e))
-                {
-                    let entry = entry.unwrap();
-                    if entry.file_type().is_dir() {
-                        continue;
-                    }
-                    if !self.recursive && entry.depth() > 1 {
-                        continue;
+    pub fn load<P: AsRef<Path>>(mut self, source: P) -> Result<Self> {
+        self.paths = match source.as_ref() {
+            s if s.is_file() => VecDeque::from([s.to_path_buf()]),
+            s if s.is_dir() => WalkDir::new(s)
+                .into_iter()
+                .filter_entry(|e| !Self::_is_hidden(e))
+                .filter_map(|entry| match entry {
+                    Err(_) => None,
+                    Ok(entry) => {
+                        if entry.file_type().is_dir() {
+                            return None;
+                        }
+                        if !self.recursive && entry.depth() > 1 {
+                            return None;
+                        }
+                        Some(entry.path().to_path_buf())
                     }
-                    paths.push_back(entry.path().to_path_buf());
-                }
-            }
+                })
+                .collect::<VecDeque<_>>(),
             // s if s.starts_with("rtsp://") || s.starts_with("rtmp://") || s.starts_with("http://")|| s.starts_with("https://") => todo!(),
             s if !s.exists() => bail!("{s:?} Not Exists"),
             _ => todo!(),
-        }
-        let n_new = paths.len();
-        self.paths.append(&mut paths);
-        println!(
-            "{CHECK_MARK} Found images x{n_new} ({} total)",
-            self.paths.len()
-        );
-        Ok(Self {
-            paths: self.paths.to_owned(),
-            batch: self.batch,
-            recursive: self.recursive,
-        })
+        };
+        println!("{CHECK_MARK} Found file x{}", self.paths.len());
+        Ok(self)
     }
 
     pub fn try_read<P: AsRef<Path>>(path: P) -> Result<DynamicImage> {

src/core/engine.rs

@@ -4,7 +4,6 @@ use human_bytes::human_bytes;
 use ndarray::{Array, IxDyn};
 use ort::{
     ExecutionProvider, Session, SessionBuilder, TensorElementType, TensorRTExecutionProvider,
-    MINOR_VERSION,
 };
 use prost::Message;
 use std::collections::HashSet;
@@ -146,7 +145,7 @@ impl OrtEngine {
 
         // summary
         println!(
-            "{CHECK_MARK} ORT: 1.{MINOR_VERSION}.x | Opset: {} | EP: {:?} | Dtype: {:?} | Parameters: {}",
+            "{CHECK_MARK} ONNX | OpSet: {} | EP: {:?} | DType: {:?} | Params: {}",
             model_proto.opset_import[0].version,
             device,
             inputs_attrs.dtypes,

src/lib.rs

@@ -1,4 +1,4 @@
-//! A Rust library integrated with ONNXRuntime, providing a collection of Computer Vision and Vision-Language models.
+//! A Rust library integrated with ONNXRuntime, providing a collection of **Computer Vision** and **Vision-Language** models.
 //!
 //! # Supported Models
 //!