Skip to content

Commit a69848d

Browse files
jamjamjonjamjamjon
committed
Add RTMPose, RTMW, DWPose models
1 parent 70aeae9 commit a69848d

25 files changed

Lines changed: 1436 additions & 9 deletions

File tree

‎README.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,10 @@
2424
## 🚀 Quick Start
2525
```bash
2626
# CPU
27-
cargo run -r --example yolo # YOLOv8-n detect by default
27+
cargo run -r --example yolo -- --task detect --ver 8 --scale n --dtype fp16 # q8, q4, q4f16
2828

2929
# NVIDIA CUDA
30-
cargo run -r -F cuda --example yolo -- --device cuda:0
30+
cargo run -r -F cuda --example yolo -- --device cuda:0 # YOLOv8-n detect by default
3131

3232
# NVIDIA TensorRT
3333
cargo run -r -F tensorrt --example yolo -- --device tensorrt:0
@@ -79,6 +79,9 @@ usls = "latest-version"
7979
| [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO) | Object Detection | [demo](examples/picodet-layout) |
8080
| [D-FINE](https://github.com/manhbd-22022602/D-FINE) | Object Detection | [demo](examples/d-fine) |
8181
| [DEIM](https://github.com/ShihuaHuang95/DEIM) | Object Detection | [demo](examples/deim) |
82+
| [RTMPose](https://github.com/open-mmlab/mmpose/tree/dev-1.x/projects/rtmpose) | Keypoint Detection | [demo](examples/rtmpose) |
83+
| [DWPose](https://github.com/IDEA-Research/DWPose) | Keypoint Detection | [demo](examples/dwpose) |
84+
| [RTMW](https://arxiv.org/abs/2407.08634) | Keypoint Detection | [demo](examples/rtmw) |
8285
| [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo) | Keypoint Detection | [demo](examples/rtmo) |
8386
| [SAM](https://github.com/facebookresearch/segment-anything) | Segment Anything | [demo](examples/sam) |
8487
| [SAM2](https://github.com/facebookresearch/segment-anything-2) | Segment Anything | [demo](examples/sam) |

‎examples/dwpose/README.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
## Quick Start
2+
3+
```shell
4+
cargo run -r --example dwpose
5+
```
6+
7+
## Results
8+
9+
![](https://github.com/jamjamjon/assets/releases/download/rtmpose/demo-dwpose.jpg)

‎examples/dwpose/main.rs‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
use anyhow::Result;
2+
use usls::{
3+
models::{RTMPose, YOLO},
4+
Annotator, Config, DataLoader, Scale, Style, SKELETON_COCO_65, SKELETON_COLOR_COCO_65,
5+
};
6+
7+
#[derive(argh::FromArgs)]
8+
/// Example
9+
struct Args {
10+
/// source: image, image folder, video stream
11+
#[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
12+
source: String,
13+
14+
/// device: cuda:0, cpu:0, ...
15+
#[argh(option, default = "String::from(\"cpu:0\")")]
16+
device: String,
17+
18+
/// scale: t, s, m, l
19+
#[argh(option, default = "String::from(\"m\")")]
20+
scale: String,
21+
22+
/// dtype: fp16, q8, q4, q4f16, ...
23+
#[argh(option, default = "String::from(\"auto\")")]
24+
dtype: String,
25+
}
26+
27+
fn main() -> Result<()> {
28+
tracing_subscriber::fmt()
29+
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
30+
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
31+
.init();
32+
let args: Args = argh::from_env();
33+
34+
// build YOLOv8
35+
let yolo_config = Config::yolo_detect()
36+
.with_scale(Scale::N)
37+
.with_version(8.into())
38+
.with_model_device(args.device.parse()?)
39+
.retain_classes(&[0]) // keep person class only
40+
.with_class_confs(&[0.5])
41+
.commit()?;
42+
let mut yolo = YOLO::new(yolo_config)?;
43+
44+
// build RTMPose
45+
let config = match args.scale.as_str() {
46+
"t" => Config::dwpose_133_t(),
47+
"s" => Config::dwpose_133_s(),
48+
"m" => Config::dwpose_133_m(),
49+
"l" => Config::dwpose_133_l(),
50+
_ => todo!(),
51+
}
52+
.with_model_dtype(args.dtype.parse()?)
53+
.with_model_device(args.device.parse()?)
54+
.commit()?;
55+
let mut rtmpose = RTMPose::new(config)?;
56+
57+
// build annotator
58+
let annotator = Annotator::default()
59+
.with_hbb_style(Style::hbb().with_draw_fill(true))
60+
.with_keypoint_style(
61+
Style::keypoint()
62+
.with_radius(2)
63+
.with_skeleton((SKELETON_COCO_65, SKELETON_COLOR_COCO_65).into())
64+
.show_id(false)
65+
.show_confidence(false)
66+
.show_name(false),
67+
);
68+
69+
// build dataloader
70+
let dl = DataLoader::new(&args.source)?.with_batch(1).build()?;
71+
72+
// iterate
73+
for xs in &dl {
74+
// YOLO infer
75+
let ys_det = yolo.forward(&xs)?;
76+
77+
// RTMPose infer
78+
for (x, y_det) in xs.iter().zip(ys_det.iter()) {
79+
let y = rtmpose.forward(x, y_det.hbbs())?;
80+
81+
// Annotate
82+
annotator.annotate(x, &y)?.save(format!(
83+
"{}.jpg",
84+
usls::Dir::Current
85+
.base_dir_with_subs(&["runs", rtmpose.spec()])?
86+
.join(usls::timestamp(None))
87+
.display(),
88+
))?
89+
}
90+
}
91+
92+
// summary
93+
yolo.summary();
94+
rtmpose.summary();
95+
96+
Ok(())
97+
}

‎examples/rtmo/main.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ struct Args {
1212
#[argh(option, default = "String::from(\"cpu:0\")")]
1313
device: String,
1414

15-
/// scale: s, m, l
15+
/// scale: t, s, m, l
1616
#[argh(option, default = "String::from(\"t\")")]
1717
scale: String,
1818
}
@@ -52,7 +52,7 @@ fn main() -> Result<()> {
5252
.with_skeleton(SKELETON_COCO_19.into())
5353
.show_confidence(false)
5454
.show_id(true)
55-
.show_name(false),
55+
.show_name(true),
5656
);
5757
for (x, y) in xs.iter().zip(ys.iter()) {
5858
annotator.annotate(x, y)?.save(format!(

‎examples/rtmpose/README.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
## Quick Start
2+
3+
```shell
4+
cargo run -r --example rtmpose -- --is-coco true
5+
```
6+
7+
## Results
8+
9+
| rtmpose-17 | rtmpose-26 |
10+
| :---------------: | :------------------------: |
11+
| ![](https://github.com/jamjamjon/assets/releases/download/rtmpose/demo-rtmpose-17.jpg) | ![](https://github.com/jamjamjon/assets/releases/download/rtmpose/demo-rtmpose-26.jpg)|

‎examples/rtmpose/main.rs‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
use anyhow::Result;
2+
use usls::{
3+
models::{RTMPose, YOLO},
4+
Annotator, Config, DataLoader, Scale, Style, SKELETON_COCO_19, SKELETON_COLOR_COCO_19,
5+
SKELETON_COLOR_HALPE_27, SKELETON_HALPE_27,
6+
};
7+
8+
#[derive(argh::FromArgs)]
9+
/// Example
10+
struct Args {
11+
/// source: image, image folder, video stream
12+
#[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
13+
source: String,
14+
15+
/// device: cuda:0, cpu:0, ...
16+
#[argh(option, default = "String::from(\"cpu:0\")")]
17+
device: String,
18+
19+
/// scale: t, s, m, l, x
20+
#[argh(option, default = "String::from(\"t\")")]
21+
scale: String,
22+
23+
/// dtype: fp16, q8, q4, q4f16, ...
24+
#[argh(option, default = "String::from(\"auto\")")]
25+
dtype: String,
26+
27+
/// is coco 17 keypoints or halpe 26 keypoints
28+
#[argh(option, default = "true")]
29+
is_coco: bool,
30+
}
31+
32+
fn main() -> Result<()> {
33+
tracing_subscriber::fmt()
34+
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
35+
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
36+
.init();
37+
let args: Args = argh::from_env();
38+
39+
// build YOLOv8
40+
let yolo_config = Config::yolo_detect()
41+
.with_scale(Scale::N)
42+
.with_version(8.into())
43+
.with_model_device(args.device.parse()?)
44+
.retain_classes(&[0]) // keep person class only
45+
.with_class_confs(&[0.5])
46+
.commit()?;
47+
let mut yolo = YOLO::new(yolo_config)?;
48+
49+
// build RTMPose
50+
let config = match args.scale.as_str() {
51+
"t" => match args.is_coco {
52+
true => Config::rtmpose_17_t(),
53+
false => Config::rtmpose_26_t(),
54+
},
55+
"s" => match args.is_coco {
56+
true => Config::rtmpose_17_s(),
57+
false => Config::rtmpose_26_s(),
58+
},
59+
"m" => match args.is_coco {
60+
true => Config::rtmpose_17_m(),
61+
false => Config::rtmpose_26_m(),
62+
},
63+
"l" => match args.is_coco {
64+
true => Config::rtmpose_17_l(),
65+
false => Config::rtmpose_26_l(),
66+
},
67+
"x" => match args.is_coco {
68+
true => Config::rtmpose_17_x(),
69+
false => Config::rtmpose_26_x(),
70+
},
71+
_ => todo!(),
72+
}
73+
.with_model_dtype(args.dtype.parse()?)
74+
.with_model_device(args.device.parse()?)
75+
.commit()?;
76+
let mut rtmpose = RTMPose::new(config)?;
77+
78+
// build annotator
79+
let annotator = Annotator::default()
80+
.with_hbb_style(Style::hbb().with_draw_fill(true))
81+
.with_keypoint_style(
82+
Style::keypoint()
83+
.with_radius(4)
84+
.with_skeleton(if args.is_coco {
85+
(SKELETON_COCO_19, SKELETON_COLOR_COCO_19).into()
86+
} else {
87+
(SKELETON_HALPE_27, SKELETON_COLOR_HALPE_27).into()
88+
})
89+
.show_id(false)
90+
.show_confidence(false)
91+
.show_name(false),
92+
);
93+
94+
// build dataloader
95+
let dl = DataLoader::new(&args.source)?.with_batch(1).build()?;
96+
97+
// iterate
98+
for xs in &dl {
99+
// YOLO infer
100+
let ys_det = yolo.forward(&xs)?;
101+
102+
// RTMPose infer
103+
for (x, y_det) in xs.iter().zip(ys_det.iter()) {
104+
let y = rtmpose.forward(x, y_det.hbbs())?;
105+
106+
// Annotate
107+
annotator.annotate(x, &y)?.save(format!(
108+
"{}.jpg",
109+
usls::Dir::Current
110+
.base_dir_with_subs(&["runs", rtmpose.spec()])?
111+
.join(usls::timestamp(None))
112+
.display(),
113+
))?
114+
}
115+
}
116+
117+
// summary
118+
yolo.summary();
119+
rtmpose.summary();
120+
121+
Ok(())
122+
}

‎examples/rtmw/README.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
## Quick Start
2+
3+
```shell
4+
cargo run -r --example rtmw
5+
```
6+
7+
## Result
8+
9+
![](https://github.com/jamjamjon/assets/releases/download/rtmpose/demo-rtmw.jpg)

‎examples/rtmw/main.rs‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
use anyhow::Result;
2+
use usls::{
3+
models::{RTMPose, YOLO},
4+
Annotator, Config, DataLoader, Scale, Style, SKELETON_COCO_65, SKELETON_COLOR_COCO_65,
5+
};
6+
7+
#[derive(argh::FromArgs)]
8+
/// Example
9+
struct Args {
10+
/// source: image, image folder, video stream
11+
#[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
12+
source: String,
13+
14+
/// device: cuda:0, cpu:0, ...
15+
#[argh(option, default = "String::from(\"cpu:0\")")]
16+
device: String,
17+
18+
/// scale: m, l, x
19+
#[argh(option, default = "String::from(\"m\")")]
20+
scale: String,
21+
22+
/// dtype: fp16, q8, q4, q4f16, ...
23+
#[argh(option, default = "String::from(\"auto\")")]
24+
dtype: String,
25+
}
26+
27+
fn main() -> Result<()> {
28+
tracing_subscriber::fmt()
29+
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
30+
.with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
31+
.init();
32+
let args: Args = argh::from_env();
33+
34+
// build YOLOv8
35+
let yolo_config = Config::yolo_detect()
36+
.with_scale(Scale::N)
37+
.with_version(8.into())
38+
.with_model_device(args.device.parse()?)
39+
.retain_classes(&[0]) // keep person class only
40+
.with_class_confs(&[0.5])
41+
.commit()?;
42+
let mut yolo = YOLO::new(yolo_config)?;
43+
44+
// build RTMPose
45+
let config = match args.scale.as_str() {
46+
"m" => Config::rtmw_133_m(),
47+
"l" => Config::rtmw_133_l(),
48+
"x" => Config::rtmw_133_x(),
49+
_ => todo!(),
50+
}
51+
.with_model_dtype(args.dtype.parse()?)
52+
.with_model_device(args.device.parse()?)
53+
.with_model_file("/HDD/zhangj/PROJECTS/e/mmml/rtmw-133-m-q4.onnx")
54+
.commit()?;
55+
let mut rtmpose = RTMPose::new(config)?;
56+
57+
// build annotator
58+
let annotator = Annotator::default()
59+
.with_hbb_style(Style::hbb().with_draw_fill(true))
60+
.with_keypoint_style(
61+
Style::keypoint()
62+
.with_radius(2)
63+
.with_skeleton((SKELETON_COCO_65, SKELETON_COLOR_COCO_65).into())
64+
.show_id(false)
65+
.show_confidence(false)
66+
.show_name(false),
67+
);
68+
69+
// build dataloader
70+
let dl = DataLoader::new(&args.source)?.with_batch(1).build()?;
71+
72+
// iterate
73+
for xs in &dl {
74+
// YOLO infer
75+
let ys_det = yolo.forward(&xs)?;
76+
77+
// RTMPose infer
78+
for (x, y_det) in xs.iter().zip(ys_det.iter()) {
79+
let y = rtmpose.forward(x, y_det.hbbs())?;
80+
81+
// Annotate
82+
annotator.annotate(x, &y)?.save(format!(
83+
"{}.jpg",
84+
usls::Dir::Current
85+
.base_dir_with_subs(&["runs", rtmpose.spec()])?
86+
.join(usls::timestamp(None))
87+
.display(),
88+
))?
89+
}
90+
}
91+
92+
// summary
93+
yolo.summary();
94+
rtmpose.summary();
95+
96+
Ok(())
97+
}

0 commit comments

Comments
 (0)