Skip to content

Commit 8c9e28e

Browse files
jamjamjonjamjamjon
authored
Add RT-DETRv4 (#161)
1 parent b89da43 commit 8c9e28e

32 files changed

Lines changed: 6784 additions & 6644 deletions

‎Cargo.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
[package]
22
name = "usls"
33
edition = "2021"
4-
version = "0.1.8"
4+
version = "0.1.9"
55
rust-version = "1.85"
66
description = "A Rust library integrated with ONNXRuntime, providing a collection of ML models."
77
repository = "https://github.com/jamjamjon/usls"
@@ -14,7 +14,7 @@ exclude = ["assets/*", "examples/*", "runs/*", "benches/*", "tests/*"]
1414
[dependencies]
1515
anyhow = { version = "1" }
1616
slsl = { version = "0.0.7", features = ["rayon"], optional = true }
17-
aksr = { version = "0.0.6" }
17+
aksr = { version = "0.0.7" }
1818
ab_glyph = { version = "0.2.32" }
1919
image = { version = "0.25" }
2020
imageproc = { version = "0.25" }

‎README.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,8 @@ usls = { version = "latest-version", features = [ "cuda" ] }
121121
| [YOLOv10](https://github.com/THU-MIG/yolov10) | Object Detection | `yolo` | [demo](examples/yolo) |
122122
| [YOLOv12](https://github.com/sunsmarterjie/yolov12) | Image Classification<br />Object Detection<br />Instance Segmentation | `yolo` | [demo](examples/yolo) |
123123
| [YOLOv13](https://github.com/iMoonLab/yolov13) | Object Detection | `yolo` | [demo](examples/yolo) |
124-
| [RT-DETR](https://github.com/lyuwenyu/RT-DETR) | Object Detection | `rtdetr` | [demo](examples/rtdetr) |
124+
| [RT-DETRv1, v2](https://github.com/lyuwenyu/RT-DETR) | Object Detection | `rtdetr` | [demo](examples/rtdetr) |
125+
| [RT-DETRv4](https://github.com/RT-DETRs/RT-DETRv4) | Object Detection | `rtdetr` | [demo](examples/rtdetr) |
125126
| [RF-DETR](https://github.com/roboflow/rf-detr) | Object Detection | `rfdetr` | [demo](examples/rfdetr) |
126127
| [PP-PicoDet](https://github.com/PaddlePaddle/PaddleDetection/tree/release/2.8/configs/picodet) | Object Detection | `picodet` | [demo](examples/picodet-layout) |
127128
| [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO) | Object Detection | `picodet` | [demo](examples/picodet-layout) |

‎examples/rfdetr/main.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,6 @@ fn main() -> Result<()> {
5151

5252
// run
5353
let ys = model.forward(&xs)?;
54-
55-
// extract bboxes
5654
println!("{:?}", ys);
5755

5856
// annotate

‎examples/rtdetr/README.md‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,7 @@
11
## Quick Start
22

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

‎examples/rtdetr/main.rs‎

Lines changed: 85 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,95 @@
11
use anyhow::Result;
2-
use usls::{models::RTDETR, Annotator, Config, DataLoader};
2+
use usls::{models::RTDETR, Annotator, Config, DataLoader, Scale, Version};
3+
4+
#[derive(argh::FromArgs)]
5+
/// Example
6+
struct Args {
7+
/// device
8+
#[argh(option, default = "String::from(\"cpu\")")]
9+
device: String,
10+
11+
/// scale
12+
#[argh(option, default = "String::from(\"s\")")]
13+
scale: String,
14+
15+
/// dtype
16+
#[argh(option, default = "String::from(\"q4f16\")")]
17+
dtype: String,
18+
19+
/// model
20+
#[argh(option)]
21+
model: Option<String>,
22+
23+
/// version
24+
#[argh(option, default = "1.0")]
25+
ver: f32,
26+
27+
/// confidences
28+
#[argh(option)]
29+
confs: Vec<f32>,
30+
}
331

432
fn main() -> Result<()> {
533
tracing_subscriber::fmt()
634
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
735
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
836
.init();
37+
let args: Args = argh::from_env();
938

1039
// config
11-
let config = Config::rtdetr_v2_s_coco().commit()?;
12-
// rtdetr_v1_r18vd_coco()
13-
// rtdetr_v2_ms_coco()
14-
// rtdetr_v2_m_coco()
15-
// rtdetr_v2_l_coco()
16-
// rtdetr_v2_x_coco()
40+
let config = match args.model {
41+
Some(model) => Config::rtdetr().with_model_file(&model),
42+
None => {
43+
match args.ver.try_into()? {
44+
Version(1, 0, _) => match args.scale.parse()? {
45+
Scale::Named(ref name) if name == "r18" => Config::rtdetr_v1_r18(),
46+
Scale::Named(ref name) if name == "r18-obj365" => Config::rtdetr_v1_r18_obj365(),
47+
Scale::Named(ref name) if name == "r34" => Config::rtdetr_v1_r34(),
48+
Scale::Named(ref name) if name == "r50" => Config::rtdetr_v1_r50(),
49+
Scale::Named(ref name) if name == "r50-obj365" => Config::rtdetr_v1_r50_obj365(),
50+
Scale::Named(ref name) if name == "r101" => Config::rtdetr_v1_r101(),
51+
Scale::Named(ref name) if name == "r101-obj365" => Config::rtdetr_v1_r101_obj365(),
52+
_ => unimplemented!(
53+
"Unsupported model scale: {:?} for RT-DETRv1. Try r18, r18-obj365, r34, r50, r50-obj365, r101, r101-obj365.",
54+
args.scale,
55+
),
56+
},
57+
Version(2, 0, _) => match args.scale.parse()?{
58+
Scale::S => Config::rtdetr_v2_s(),
59+
Scale::M => Config::rtdetr_v2_m(),
60+
Scale::Named(ref name) if name == "ms" => Config::rtdetr_v2_ms(),
61+
Scale::L => Config::rtdetr_v2_l(),
62+
Scale::X => Config::rtdetr_v2_x(),
63+
_ => unimplemented!(
64+
"Unsupported model scale: {:?} for RT-DETRv2. Try s, m, ms, l, x.",
65+
args.scale,
66+
),
67+
},
68+
Version(4, 0, _) => match args.scale.parse()?{
69+
Scale::S => Config::rtdetr_v4_s(),
70+
Scale::M => Config::rtdetr_v4_m(),
71+
Scale::L => Config::rtdetr_v4_l(),
72+
Scale::X => Config::rtdetr_v4_x(),
73+
_ => unimplemented!(
74+
"Unsupported model scale: {:?} for RT-DETRv4. Try s, m, l, x.",
75+
args.scale,
76+
),
77+
},
78+
_ => unimplemented!(
79+
"Unsupported model version: {:?}. Try v1, v2, v4 for RT-DETR.",
80+
args.ver
81+
),
82+
}
83+
}
84+
}
85+
.with_dtype_all(args.dtype.parse()?)
86+
.with_device_all(args.device.parse()?)
87+
.with_class_confs(if args.confs.is_empty() {
88+
&[0.35]
89+
} else {
90+
&args.confs
91+
})
92+
.commit()?;
1793
let mut model = RTDETR::new(config)?;
1894

1995
// load
@@ -34,6 +110,8 @@ fn main() -> Result<()> {
34110
.display(),
35111
))?;
36112
}
113+
114+
// summary
37115
usls::perf(false);
38116

39117
Ok(())

‎src/core/mod.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ pub mod image;
2929
mod logits_sampler;
3030
mod media;
3131
mod min_opt_max;
32-
mod names;
3332
#[cfg(any(feature = "ort-download-binaries", feature = "ort-load-dynamic"))]
3433
#[allow(clippy::all)]
3534
pub(crate) mod onnx;
@@ -74,7 +73,6 @@ pub use image::*;
7473
pub use logits_sampler::LogitsSampler;
7574
pub use media::*;
7675
pub use min_opt_max::MinOptMax;
77-
pub use names::*;
7876
pub use ops::*;
7977
pub use ort_config::ORTConfig;
8078
pub use processor::*;

0 commit comments

Comments
 (0)