From 557ec87dd76d621b2cf7dd5b85ca61bf2439e439 Mon Sep 17 00:00:00 2001 From: jamjamjon Date: Thu, 25 Jul 2024 21:40:08 +0800 Subject: [PATCH] update --- Cargo.toml | 13 ++++- README.md | 5 +- benches/yolo.rs | 105 +++++++++++++++++++++++++++++++++++++++++ examples/yolo/main.rs | 2 +- src/core/dataloader.rs | 53 +++++++++------------ src/core/engine.rs | 3 +- src/lib.rs | 2 +- 7 files changed, 145 insertions(+), 38 deletions(-) create mode 100644 benches/yolo.rs diff --git a/Cargo.toml b/Cargo.toml index 20867ba..ad1b662 100644 --- a/Cargo.toml +++ b/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 diff --git a/README.md b/README.md index 5b5f38f..6bc06c5 100644 --- a/README.md +++ b/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. diff --git a/benches/yolo.rs b/benches/yolo.rs new file mode 100644 index 0000000..7398748 --- /dev/null +++ b/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); diff --git a/examples/yolo/main.rs b/examples/yolo/main.rs index df64b24..1879a1f 100644 --- a/examples/yolo/main.rs +++ b/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"); diff --git a/src/core/dataloader.rs b/src/core/dataloader.rs index e7bb444..bbce77d 100644 --- a/src/core/dataloader.rs +++ b/src/core/dataloader.rs @@ -53,42 +53,31 @@ impl Default for DataLoader { } impl DataLoader { - pub fn load>(&mut self, source: P) -> Result { - 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>(mut self, source: P) -> Result { + 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::>(), // 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>(path: P) -> Result { diff --git a/src/core/engine.rs b/src/core/engine.rs index ee4dc9d..3f5c151 100644 --- a/src/core/engine.rs +++ b/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, diff --git a/src/lib.rs b/src/lib.rs index 10e3158..e4da5e1 100644 --- a/src/lib.rs +++ b/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 //!