Add RT-DETRv4

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "usls"
 edition = "2021"
-version = "0.1.8"
+version = "0.1.9"
 rust-version = "1.85"
 description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models."
 repository = "https://github.com/jamjamjon/usls"
@@ -14,7 +14,7 @@ exclude = ["assets/*", "examples/*", "runs/*", "benches/*", "tests/*"]
 [dependencies]
 anyhow = { version = "1" }
 slsl = { version = "0.0.7", features = ["rayon"], optional = true }
-aksr = { version = "0.0.6" }
+aksr = { version = "0.0.7" }
 ab_glyph = { version = "0.2.32" }
 image = { version = "0.25" }
 imageproc = { version = "0.25" }

README.md

@@ -121,7 +121,8 @@ usls = { version = "latest-version", features = [ "cuda" ] }
 | [YOLOv10](https://github.com/THU-MIG/yolov10) | Object Detection | `yolo` | [demo](examples/yolo) |
 | [YOLOv12](https://github.com/sunsmarterjie/yolov12) | Image Classification<br />Object Detection<br />Instance Segmentation | `yolo` | [demo](examples/yolo) |
 | [YOLOv13](https://github.com/iMoonLab/yolov13) | Object Detection | `yolo` | [demo](examples/yolo) |
-| [RT-DETR](https://github.com/lyuwenyu/RT-DETR) | Object Detection | `rtdetr` | [demo](examples/rtdetr) |
+| [RT-DETRv1, v2](https://github.com/lyuwenyu/RT-DETR) | Object Detection | `rtdetr` | [demo](examples/rtdetr) |
+| [RT-DETRv4](https://github.com/RT-DETRs/RT-DETRv4) | Object Detection | `rtdetr` | [demo](examples/rtdetr) |
 | [RF-DETR](https://github.com/roboflow/rf-detr) | Object Detection | `rfdetr` | [demo](examples/rfdetr) |
 | [PP-PicoDet](https://github.com/PaddlePaddle/PaddleDetection/tree/release/2.8/configs/picodet) | Object Detection | `picodet` | [demo](examples/picodet-layout) |
 | [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO) | Object Detection | `picodet` | [demo](examples/picodet-layout) |

examples/rfdetr/main.rs

@@ -51,8 +51,6 @@ fn main() -> Result<()> {
 
     // run
     let ys = model.forward(&xs)?;
-
-    // extract bboxes
     println!("{:?}", ys);
 
     // annotate

examples/rtdetr/README.md

@@ -1,17 +1,7 @@
 ## Quick Start
 
 ```shell
-cargo run -r -F rtdetr --example rtdetr
-```
-
-## Results
-
-```
-[Bboxes]: Found 5 objects
-0: Bbox { xyxy: [47.969677, 397.81808, 246.22426, 904.8823], class_id: 0, name: Some("person"), confidence: 0.94432133 }
-1: Bbox { xyxy: [668.0796, 399.28854, 810.3779, 880.7412], class_id: 0, name: Some("person"), confidence: 0.93386495 }
-2: Bbox { xyxy: [20.852705, 229.30482, 807.43494, 729.51196], class_id: 5, name: Some("bus"), confidence: 0.9319465 }
-3: Bbox { xyxy: [223.28226, 405.37265, 343.92603, 859.50366], class_id: 0, name: Some("person"), confidence: 0.9130827 }
-4: Bbox { xyxy: [0.0, 552.6165, 65.99908, 868.00525], class_id: 0, name: Some("person"), confidence: 0.7910869 }
-
+cargo run -F rtdetr --example rtdetr -- --scale r18 --dtype fp16 --device cuda:0 --confs 0.5
+cargo run -F rtdetr --example rtdetr -- --scale s --ver 2.0
+cargo run -F rtdetr --example rtdetr -- --scale m --ver 4.0
 ```

examples/rtdetr/main.rs

@@ -1,19 +1,95 @@
 use anyhow::Result;
-use usls::{models::RTDETR, Annotator, Config, DataLoader};
+use usls::{models::RTDETR, Annotator, Config, DataLoader, Scale, Version};
+
+#[derive(argh::FromArgs)]
+/// Example
+struct Args {
+    /// device
+    #[argh(option, default = "String::from(\"cpu\")")]
+    device: String,
+
+    /// scale
+    #[argh(option, default = "String::from(\"s\")")]
+    scale: String,
+
+    /// dtype
+    #[argh(option, default = "String::from(\"q4f16\")")]
+    dtype: String,
+
+    /// model
+    #[argh(option)]
+    model: Option<String>,
+
+    /// version
+    #[argh(option, default = "1.0")]
+    ver: f32,
+
+    /// confidences
+    #[argh(option)]
+    confs: Vec<f32>,
+}
 
 fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
         .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
         .init();
+    let args: Args = argh::from_env();
 
     // config
-    let config = Config::rtdetr_v2_s_coco().commit()?;
-    // rtdetr_v1_r18vd_coco()
-    // rtdetr_v2_ms_coco()
-    // rtdetr_v2_m_coco()
-    // rtdetr_v2_l_coco()
-    // rtdetr_v2_x_coco()
+    let config = match args.model {
+        Some(model) => Config::rtdetr().with_model_file(&model),
+        None => {
+            match args.ver.try_into()? {
+                Version(1, 0, _) => match args.scale.parse()? {
+                    Scale::Named(ref name) if name == "r18" => Config::rtdetr_v1_r18(),
+                    Scale::Named(ref name) if name == "r18-obj365" => Config::rtdetr_v1_r18_obj365(),
+                    Scale::Named(ref name) if name == "r34" => Config::rtdetr_v1_r34(),
+                    Scale::Named(ref name) if name == "r50" => Config::rtdetr_v1_r50(),
+                    Scale::Named(ref name) if name == "r50-obj365" => Config::rtdetr_v1_r50_obj365(),
+                    Scale::Named(ref name) if name == "r101" => Config::rtdetr_v1_r101(),
+                    Scale::Named(ref name) if name == "r101-obj365" => Config::rtdetr_v1_r101_obj365(),
+                    _ => unimplemented!(
+                        "Unsupported model scale: {:?} for RT-DETRv1. Try r18, r18-obj365, r34, r50, r50-obj365, r101, r101-obj365.",
+                        args.scale,
+                    ),
+                },
+                Version(2, 0, _) => match args.scale.parse()?{
+                    Scale::S => Config::rtdetr_v2_s(),
+                    Scale::M => Config::rtdetr_v2_m(),
+                    Scale::Named(ref name) if name == "ms" => Config::rtdetr_v2_ms(),
+                    Scale::L => Config::rtdetr_v2_l(),
+                    Scale::X => Config::rtdetr_v2_x(),
+                    _ => unimplemented!(
+                        "Unsupported model scale: {:?} for RT-DETRv2. Try s, m, ms, l, x.",
+                        args.scale,
+                    ),
+                },
+                Version(4, 0, _) => match args.scale.parse()?{
+                    Scale::S => Config::rtdetr_v4_s(),
+                    Scale::M => Config::rtdetr_v4_m(),
+                    Scale::L => Config::rtdetr_v4_l(),
+                    Scale::X => Config::rtdetr_v4_x(),
+                    _ => unimplemented!(
+                        "Unsupported model scale: {:?} for RT-DETRv4. Try s, m, l, x.",
+                        args.scale,
+                    ),
+                },
+                _ => unimplemented!(
+                    "Unsupported model version: {:?}. Try v1, v2, v4 for RT-DETR.",
+                    args.ver
+                ),
+            }
+        }
+    }
+    .with_dtype_all(args.dtype.parse()?)
+    .with_device_all(args.device.parse()?)
+    .with_class_confs(if args.confs.is_empty() {
+            &[0.35]
+        } else {
+            &args.confs
+        })
+    .commit()?;
     let mut model = RTDETR::new(config)?;
 
     // load
@@ -34,6 +110,8 @@ fn main() -> Result<()> {
                 .display(),
         ))?;
     }
+
+    // summary
     usls::perf(false);
 
     Ok(())

src/core/mod.rs

@@ -29,7 +29,6 @@ pub mod image;
 mod logits_sampler;
 mod media;
 mod min_opt_max;
-mod names;
 #[cfg(any(feature = "ort-download-binaries", feature = "ort-load-dynamic"))]
 #[allow(clippy::all)]
 pub(crate) mod onnx;
@@ -74,7 +73,6 @@ pub use image::*;
 pub use logits_sampler::LogitsSampler;
 pub use media::*;
 pub use min_opt_max::MinOptMax;
-pub use names::*;
 pub use ops::*;
 pub use ort_config::ORTConfig;
 pub use processor::*;