feat: support HRNet DarkPose pose estimation (w32/w48, COCO body & wholebody)

examples/pose-estimation/hrnet.rs

@@ -0,0 +1,86 @@
+use anyhow::Result;
+use clap::Args;
+use usls::{Config, DType, Device};
+
+#[derive(Args, Debug)]
+pub struct HrnetArgs {
+    /// Backbone width: w32 or w48
+    #[arg(long, default_value = "w48")]
+    pub width: String,
+
+    /// 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). WholeBody always uses 384x288.
+    #[arg(long, default_value_t = true, action = clap::ArgAction::Set)]
+    pub hires: bool,
+
+    /// Optional local model file path (overrides the built-in selection)
+    #[arg(long)]
+    pub model: Option<String>,
+
+    /// 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: &HrnetArgs) -> Result<Config> {
+    let mut config = match (args.width.as_str(), args.is_coco) {
+        ("w32", true) => {
+            if args.hires {
+                Config::hrnet_w32_17_384()
+            } else {
+                Config::hrnet_w32_17()
+            }
+        }
+        ("w48", true) => {
+            if args.hires {
+                Config::hrnet_w48_17_384()
+            } else {
+                Config::hrnet_w48_17()
+            }
+        }
+        ("w32", false) => Config::hrnet_w32_133(),
+        ("w48", false) => Config::hrnet_w48_133(),
+        (w, _) => anyhow::bail!("Unsupported HRNet width: {w} (expected w32 or w48)"),
+    };
+
+    // Allow overriding with a local file (e.g. the sample end2end.hrnet_w48.onnx)
+    if let Some(model) = &args.model {
+        config = config.with_model_file(model);
+    }
+
+    let config = config
+        .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)
+}

examples/pose-estimation/main.rs

@@ -1,11 +1,12 @@
 use anyhow::Result;
 use clap::{Parser, Subcommand};
 use usls::{
-    models::{DWPose, RTMPose, RTMO, YOLO},
+    models::{DWPose, HRNet, RTMPose, RTMO, YOLO},
     Annotator, Config, DataLoader, Model, Scale, Source, Y,
 };
 
 mod dwpose;
+mod hrnet;
 mod rtmo;
 mod rtmpose;
 mod rtmw;
@@ -29,6 +30,7 @@ struct Cli {
 #[derive(Subcommand)]
 enum Commands {
     Dwpose(dwpose::DwposeArgs),
+    Hrnet(hrnet::HrnetArgs),
     Rtmo(rtmo::RtmoArgs),
     Rtmpose(rtmpose::RtmposeArgs),
     Rtmw(rtmw::RtmwArgs),
@@ -133,6 +135,34 @@ fn main() -> Result<()> {
                     ),
             )
         }
+        Commands::Hrnet(args) => {
+            let yolo_config = yolo_config
+                .with_model_device(args.device)
+                .with_image_processor_device(args.processor_device)
+                .commit()?;
+            let pose_config = hrnet::config(args)?.commit()?;
+            let is_coco = args.is_coco;
+            run_with_detector::<YOLO, HRNet>(
+                yolo_config,
+                pose_config,
+                &cli.source,
+                "hrnet",
+                &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::Rtmw(args) => {
             let yolo_config = yolo_config
                 .with_model_device(args.device)

src/models/vision/hrnet/config.rs

@@ -0,0 +1,116 @@
+use crate::{NAMES_COCO_KEYPOINTS_133, NAMES_COCO_KEYPOINTS_17};
+
+///
+/// > # HRNet: Deep High-Resolution Representation Learning for Human Pose Estimation
+/// >
+/// > Top-down heatmap-based pose estimator that maintains high-resolution
+/// > representations through the whole network.
+/// >
+/// > # Paper & Code
+/// >
+/// > - **Paper**: [Deep High-Resolution Representation Learning for Human Pose Estimation](https://arxiv.org/abs/1902.09212)
+/// > - **GitHub**: [open-mmlab/mmpose](https://github.com/open-mmlab/mmpose/tree/main/configs/body_2d_keypoint/topdown_heatmap)
+/// >
+/// > # Model Variants
+/// >
+/// > - **hrnet-w32 / hrnet-w48**: backbone widths (32 / 48 channels)
+/// > - **17 keypoints**: COCO body pose estimation
+/// > - **133 keypoints**: COCO-WholeBody pose estimation
+/// > - **256x192 / 384x288**: supported input resolutions
+/// >
+/// > # Implemented Features / Tasks
+/// >
+/// > - [X] **Body Pose Estimation**: 17-keypoint COCO pose estimation
+/// > - [X] **Whole-body Pose Estimation**: 133-keypoint COCO-WholeBody pose estimation
+/// > - [X] **Multiple Backbones**: w32 / w48
+/// > - [X] **Multiple Resolutions**: 256x192 and 384x288 inputs
+/// >
+/// Model configuration for `HRNet`
+///
+impl crate::Config {
+    /// Base configuration for HRNet models (256x192 input)
+    pub fn hrnet() -> Self {
+        Self::default()
+            .with_name("hrnet")
+            .with_model_ixx(0, 0, 1)
+            .with_model_ixx(0, 1, 3)
+            .with_model_ixx(0, 2, 256)
+            .with_model_ixx(0, 3, 192)
+            .with_image_mean([123.675, 116.28, 103.53])
+            .with_image_std([58.395, 57.12, 57.375])
+            .with_normalize(false) // matters!
+            .with_keypoint_confs(&[0.35])
+    }
+
+    /// HRNet for 384x288 input
+    pub fn hrnet_384() -> Self {
+        Self::hrnet()
+            .with_model_ixx(0, 2, 384)
+            .with_model_ixx(0, 3, 288)
+    }
+
+    /// Base configuration for 17-keypoint COCO body pose estimation
+    pub fn hrnet_17() -> Self {
+        Self::hrnet()
+            .with_nk(17)
+            .with_keypoint_names(&NAMES_COCO_KEYPOINTS_17)
+    }
+
+    /// Base configuration for 133-keypoint COCO-WholeBody pose estimation (256x192 input)
+    pub fn hrnet_133() -> Self {
+        Self::hrnet()
+            .with_nk(133)
+            .with_keypoint_names(&NAMES_COCO_KEYPOINTS_133)
+    }
+
+    /// HRNet-w32, 17-keypoint COCO body, 256x192 (DarkPose)
+    pub fn hrnet_w32_17() -> Self {
+        Self::hrnet_17().with_model_file(
+            "https://github.com/wep21/assets/releases/download/hrnet/hrnet-w32-coco-256x192-dark.onnx",
+        )
+    }
+
+    /// HRNet-w32, 17-keypoint COCO body, 384x288 (DarkPose)
+    pub fn hrnet_w32_17_384() -> Self {
+        Self::hrnet_17()
+            .with_model_ixx(0, 2, 384)
+            .with_model_ixx(0, 3, 288)
+            .with_model_file(
+                "https://github.com/wep21/assets/releases/download/hrnet/hrnet-w32-coco-384x288-dark.onnx",
+            )
+    }
+
+    /// HRNet-w48, 17-keypoint COCO body, 256x192 (DarkPose)
+    pub fn hrnet_w48_17() -> Self {
+        Self::hrnet_17().with_model_file(
+            "https://github.com/wep21/assets/releases/download/hrnet/hrnet-w48-coco-256x192-dark.onnx",
+        )
+    }
+
+    /// HRNet-w48, 17-keypoint COCO body, 384x288 (DarkPose)
+    pub fn hrnet_w48_17_384() -> Self {
+        Self::hrnet_17()
+            .with_model_ixx(0, 2, 384)
+            .with_model_ixx(0, 3, 288)
+            .with_model_file(
+                "https://github.com/wep21/assets/releases/download/hrnet/hrnet-w48-coco-384x288-dark.onnx",
+            )
+    }
+
+    /// HRNet-w32, 133-keypoint COCO-WholeBody, 256x192 (DarkPose)
+    pub fn hrnet_w32_133() -> Self {
+        Self::hrnet_133().with_model_file(
+            "https://github.com/wep21/assets/releases/download/hrnet/hrnet-w32-coco-wholebody-256x192-dark.onnx",
+        )
+    }
+
+    /// HRNet-w48, 133-keypoint COCO-WholeBody, 384x288 (DarkPose)
+    pub fn hrnet_w48_133() -> Self {
+        Self::hrnet_133()
+            .with_model_ixx(0, 2, 384)
+            .with_model_ixx(0, 3, 288)
+            .with_model_file(
+                "https://github.com/wep21/assets/releases/download/hrnet/hrnet-w48-coco-wholebody-384x288-dark.onnx",
+            )
+    }
+}