From e9b86df5eeb98f09dce08ac968c20debadf75410 Mon Sep 17 00:00:00 2001 From: wep21 Date: Sun, 7 Jun 2026 15:41:12 +0900 Subject: [PATCH] feat: add cigpose Signed-off-by: wep21 --- docs/model-zoo/pose.md | 1 + examples/pose-estimation/README.md | 5 ++ examples/pose-estimation/cigpose.rs | 75 +++++++++++++++++++++++++++++ examples/pose-estimation/main.rs | 32 +++++++++++- src/models/vision/cigpose/config.rs | 73 ++++++++++++++++++++++++++++ src/models/vision/cigpose/mod.rs | 3 ++ src/models/vision/mod.rs | 4 +- 7 files changed, 191 insertions(+), 2 deletions(-) create mode 100644 examples/pose-estimation/cigpose.rs create mode 100644 src/models/vision/cigpose/config.rs create mode 100644 src/models/vision/cigpose/mod.rs diff --git a/docs/model-zoo/pose.md b/docs/model-zoo/pose.md index d617341..61248c8 100644 --- a/docs/model-zoo/pose.md +++ b/docs/model-zoo/pose.md @@ -10,6 +10,7 @@ hide: | Model | Task / Description | Demo | Dynamic Batch | TensorRT | FP32 | FP16 | Q8 | Q4f16 | BNB4 | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | [RTMPose](https://github.com/open-mmlab/mmpose/tree/dev-1.x/projects/rtmpose) | Keypoint Detection | [demo](https://github.com/jamjamjon/usls/tree/main/examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ✅ | +| CIGPose | Keypoint Detection | [demo](https://github.com/jamjamjon/usls/tree/main/examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ❌ | ❌ | ❌ | | [DWPose](https://github.com/IDEA-Research/DWPose) | Keypoint Detection | [demo](https://github.com/jamjamjon/usls/tree/main/examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ✅ | | [RTMW](https://arxiv.org/abs/2407.08634) | Keypoint Detection | [demo](https://github.com/jamjamjon/usls/tree/main/examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ✅ | | [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo) | Keypoint Detection | [demo](https://github.com/jamjamjon/usls/tree/main/examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ❌ | diff --git a/examples/pose-estimation/README.md b/examples/pose-estimation/README.md index 70ddca0..b968edf 100644 --- a/examples/pose-estimation/README.md +++ b/examples/pose-estimation/README.md @@ -19,6 +19,11 @@ cargo run -F cuda-full --example pose-estimation -- rtmw --dtype f16 --device cu cargo run -F cuda-full --example pose-estimation -- rtmpose --dtype f32 --device cuda:0 --processor-device cuda:0 ``` +### CIGPose +```bash +cargo run -F cuda-full --example pose-estimation -- cigpose --dtype f16 --device cuda:0 --processor-device cuda:0 +``` + ### DWPose ```bash cargo run -F cuda-full --example pose-estimation -- dwpose --dtype f16 --device cuda:0 --processor-device cuda:0 diff --git a/examples/pose-estimation/cigpose.rs b/examples/pose-estimation/cigpose.rs new file mode 100644 index 0000000..b2d4330 --- /dev/null +++ b/examples/pose-estimation/cigpose.rs @@ -0,0 +1,75 @@ +use anyhow::Result; +use clap::Args; +use usls::{Config, DType, Device, Scale}; + +#[derive(Args, Debug)] +pub struct CigposeArgs { + /// Scale: l, x + #[arg(long, default_value = "l")] + pub scale: Scale, + + /// Use COCO 17 keypoints (true = body) or COCO-WholeBody 133 keypoints (false) + #[arg(long, default_value_t = true, action = clap::ArgAction::Set)] + pub is_coco: bool, + + /// Use 384x288 input (true) or 256x192 (false). X scale always uses 384x288. + #[arg(long, default_value_t = true, action = clap::ArgAction::Set)] + pub hires: bool, + + /// Dtype: fp32, fp16, q4f16, etc. + #[arg(long, default_value = "fp32")] + pub dtype: DType, + + /// Device: cpu, cuda:0, mps, coreml, openvino:CPU, etc. + #[arg(long, global = true, default_value = "cpu")] + pub device: Device, + + /// Processor device (for pre/post processing) + #[arg(long, global = true, default_value = "cpu")] + pub processor_device: Device, + + /// Batch size + #[arg(long, global = true, default_value_t = 1)] + pub batch: usize, + + /// Min batch size (TensorRT) + #[arg(long, global = true, default_value_t = 1)] + pub min_batch: usize, + + /// Max batch size (TensorRT) + #[arg(long, global = true, default_value_t = 4)] + pub max_batch: usize, + + /// num dry run + #[arg(long, global = true, default_value_t = 3)] + pub num_dry_run: usize, +} + +pub fn config(args: &CigposeArgs) -> Result { + let config = match (args.scale.clone(), args.is_coco) { + (Scale::L, true) => { + if args.hires { + Config::cigpose_17_l_384() + } else { + Config::cigpose_17_l() + } + } + (Scale::L, false) => { + if args.hires { + Config::cigpose_133_l_384() + } else { + Config::cigpose_133_l() + } + } + (Scale::X, false) => Config::cigpose_133_x_384(), + (Scale::X, true) => anyhow::bail!("CIGPose x scale is only available for COCO-WholeBody"), + (scale, _) => anyhow::bail!("Unsupported CIGPose scale: {scale} (expected l or x)"), + } + .with_model_dtype(args.dtype) + .with_model_device(args.device) + .with_model_batch_size_min_opt_max(args.min_batch, args.batch, args.max_batch) + .with_model_num_dry_run(args.num_dry_run) + .with_image_processor_device(args.processor_device); + + Ok(config) +} diff --git a/examples/pose-estimation/main.rs b/examples/pose-estimation/main.rs index 538a3e3..aa4171e 100644 --- a/examples/pose-estimation/main.rs +++ b/examples/pose-estimation/main.rs @@ -1,10 +1,11 @@ use anyhow::Result; use clap::{Parser, Subcommand}; use usls::{ - models::{DWPose, HRNet, RTMPose, RTMO, YOLO}, + models::{CIGPose, DWPose, HRNet, RTMPose, RTMO, YOLO}, Annotator, Config, DataLoader, Model, Scale, Source, Y, }; +mod cigpose; mod dwpose; mod hrnet; mod rtmo; @@ -29,6 +30,7 @@ struct Cli { #[derive(Subcommand)] enum Commands { + Cigpose(cigpose::CigposeArgs), Dwpose(dwpose::DwposeArgs), Hrnet(hrnet::HrnetArgs), Rtmo(rtmo::RtmoArgs), @@ -107,6 +109,34 @@ fn main() -> Result<()> { ), ) } + Commands::Cigpose(args) => { + let yolo_config = yolo_config + .with_model_device(args.device) + .with_image_processor_device(args.processor_device) + .commit()?; + let pose_config = cigpose::config(args)?.commit()?; + let is_coco = args.is_coco; + run_with_detector::( + yolo_config, + pose_config, + &cli.source, + "cigpose", + &annotator + .with_hbb_style(usls::HbbStyle::default().with_draw_fill(true)) + .with_keypoint_style( + usls::KeypointStyle::default() + .with_radius(if is_coco { 2 } else { 1 }) + .with_skeleton(if is_coco { + (usls::SKELETON_COCO_19, usls::SKELETON_COLOR_COCO_19).into() + } else { + (usls::SKELETON_COCO_65, usls::SKELETON_COLOR_COCO_65).into() + }) + .show_id(false) + .show_confidence(false) + .show_name(false), + ), + ) + } Commands::Rtmpose(args) => { let yolo_config = yolo_config .with_model_device(args.device) diff --git a/src/models/vision/cigpose/config.rs b/src/models/vision/cigpose/config.rs new file mode 100644 index 0000000..03ce5cc --- /dev/null +++ b/src/models/vision/cigpose/config.rs @@ -0,0 +1,73 @@ +use crate::{NAMES_COCO_KEYPOINTS_133, NAMES_COCO_KEYPOINTS_17}; + +const CIGPOSE_RELEASE: &str = "https://github.com/wep21/assets/releases/download/cigpose"; + +/// +/// > # CIGPose +/// > +/// > SimCC-based human pose estimation models exported to ONNX. +/// > +/// > # Model Variants +/// > +/// > - **cigpose-17-l**: Large model for 17-keypoint COCO pose estimation +/// > - **cigpose-17-l-384**: Large model for 17-keypoint COCO pose estimation with 384x288 input +/// > - **cigpose-133-l**: Large model for 133-keypoint COCO-WholeBody pose estimation +/// > - **cigpose-133-l-384**: Large model for 133-keypoint COCO-WholeBody pose estimation with 384x288 input +/// > - **cigpose-133-x-384**: Extra large model for 133-keypoint COCO-WholeBody pose estimation with 384x288 input +/// > +/// Model configuration for `CIGPose` +/// +impl crate::Config { + /// Base configuration for CIGPose models. + pub fn cigpose() -> Self { + Self::rtmpose().with_name("cigpose") + } + + /// Base configuration for 17-keypoint COCO pose estimation. + pub fn cigpose_17() -> Self { + Self::cigpose() + .with_nk(17) + .with_keypoint_names(&NAMES_COCO_KEYPOINTS_17) + } + + /// Base configuration for 133-keypoint COCO-WholeBody pose estimation. + pub fn cigpose_133() -> Self { + Self::cigpose() + .with_nk(133) + .with_keypoint_names(&NAMES_COCO_KEYPOINTS_133) + } + + /// Large model for 17-keypoint COCO pose estimation. + pub fn cigpose_17_l() -> Self { + Self::cigpose_17().with_model_file(format!("{CIGPOSE_RELEASE}/cigpose-17-l.onnx")) + } + + /// Large model for 17-keypoint COCO pose estimation with 384x288 input. + pub fn cigpose_17_l_384() -> Self { + Self::cigpose_17() + .with_model_ixx(0, 2, 384) + .with_model_ixx(0, 3, 288) + .with_model_file(format!("{CIGPOSE_RELEASE}/cigpose-17-l-384.onnx")) + } + + /// Large model for 133-keypoint COCO-WholeBody pose estimation. + pub fn cigpose_133_l() -> Self { + Self::cigpose_133().with_model_file(format!("{CIGPOSE_RELEASE}/cigpose-133-l.onnx")) + } + + /// Large model for 133-keypoint COCO-WholeBody pose estimation with 384x288 input. + pub fn cigpose_133_l_384() -> Self { + Self::cigpose_133() + .with_model_ixx(0, 2, 384) + .with_model_ixx(0, 3, 288) + .with_model_file(format!("{CIGPOSE_RELEASE}/cigpose-133-l-384.onnx")) + } + + /// Extra large model for 133-keypoint COCO-WholeBody pose estimation with 384x288 input. + pub fn cigpose_133_x_384() -> Self { + Self::cigpose_133() + .with_model_ixx(0, 2, 384) + .with_model_ixx(0, 3, 288) + .with_model_file(format!("{CIGPOSE_RELEASE}/cigpose-133-x-384.onnx")) + } +} diff --git a/src/models/vision/cigpose/mod.rs b/src/models/vision/cigpose/mod.rs new file mode 100644 index 0000000..07883d8 --- /dev/null +++ b/src/models/vision/cigpose/mod.rs @@ -0,0 +1,3 @@ +mod config; + +pub type CIGPose = crate::RTMPose; diff --git a/src/models/vision/mod.rs b/src/models/vision/mod.rs index 82808e2..c35c720 100644 --- a/src/models/vision/mod.rs +++ b/src/models/vision/mod.rs @@ -7,7 +7,7 @@ //! - **Classification**: `beit`, `convnext`, `deit`, `fastvit`, `mobileone` (config-only, use with `ImageClassifier`) //! - **Detection**: `yolo`, `yolop`, `rtdetr`, `rfdetr`, `picodet`, `d_fine`, `deim`, `deimv2` //! - **Segmentation**: `sam`, `sam2`, `mediapipe_segmenter`, `sapiens` -//! - **Pose**: `rtmpose`, `rtmw`, `dwpose`, `rtmo`, `hrnet` +//! - **Pose**: `rtmpose`, `rtmw`, `dwpose`, `rtmo`, `hrnet`, `cigpose` //! - **Depth**: `depth_anything`, `depth_pro` //! - **Feature**: `dinov2`, `dinov3`, `clip`, `blip`, `ram` //! - **OCR**: `db`, `fast`, `linknet`, `svtr`, `slanet` @@ -44,6 +44,7 @@ mod sapiens; mod yoloe_prompt_free; // Pose Estimation +mod cigpose; mod dwpose; mod hrnet; mod rtmo; @@ -93,6 +94,7 @@ pub use apisr::*; pub use beit::*; pub use ben2::*; pub use birefnet::*; +pub use cigpose::*; pub use convnext::*; pub use d_fine::*; pub use db::*;