diff --git a/README.md b/README.md index cc20978..7eeff6c 100644 --- a/README.md +++ b/README.md @@ -310,6 +310,7 @@ cargo run -r -F openvino -F ort-load-dynamic --example yolo -- --device openvino | [DWPose](https://github.com/IDEA-Research/DWPose) | Keypoint Detection | [demo](./examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ✅ | | [RTMW](https://arxiv.org/abs/2407.08634) | Keypoint Detection | [demo](./examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ✅ | | [RTMO](https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo) | Keypoint Detection | [demo](./examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ✅ | ✅ | ❌ | +| [ECPose](https://github.com/Intellindust-AI-Lab/EdgeCrafter) | Keypoint Detection | [demo](./examples/pose-estimation) | ✅ | ❓ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/pose-estimation/README.md b/examples/pose-estimation/README.md index b968edf..a625e6e 100644 --- a/examples/pose-estimation/README.md +++ b/examples/pose-estimation/README.md @@ -9,6 +9,11 @@ This directory contains examples for human pose estimation models. cargo run -F cuda-full --example pose-estimation -- rtmo --dtype f16 --device cuda:0 --processor-device cuda:0 ``` +### ECPose +```bash +cargo run -F cuda-full --example pose-estimation -- ecpose --scale s --dtype fp16 --device cuda:0 --processor-device cuda:0 +``` + ### RTMW ```bash cargo run -F cuda-full --example pose-estimation -- rtmw --dtype f16 --device cuda:0 --processor-device cuda:0 diff --git a/examples/pose-estimation/ecpose.rs b/examples/pose-estimation/ecpose.rs new file mode 100644 index 0000000..ffb49e9 --- /dev/null +++ b/examples/pose-estimation/ecpose.rs @@ -0,0 +1,105 @@ +use anyhow::Result; +use clap::Args; +use usls::{Config, DType, Device, Scale}; + +#[derive(Args, Debug)] +pub struct EcposeArgs { + /// Scale: s, m, l, x + #[arg(long, default_value = "s")] + pub scale: Scale, + + /// Optional local model file path (overrides the built-in selection) + #[arg(long)] + pub model: Option, + + /// 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: &EcposeArgs) -> Result { + let mut config = match args.scale { + Scale::S => Config::ecpose_s(), + Scale::M => Config::ecpose_m(), + Scale::L => Config::ecpose_l(), + Scale::X => Config::ecpose_x(), + _ => anyhow::bail!( + "Unsupported ECPose scale: {} (expected s, m, l, or x)", + args.scale + ), + }; + + config = if let Some(model) = &args.model { + config.with_model_file(select_local_model_path(model, args.dtype)) + } else { + config.with_model_dtype(args.dtype) + }; + + Ok(config + .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)) +} + +fn select_local_model_path(path: &str, dtype: DType) -> String { + if matches!(dtype, DType::Auto | DType::Fp32) { + return path.to_string(); + } + + let Some(stem) = path.strip_suffix(".onnx") else { + return path.to_string(); + }; + + for suffix in [ + "-fp16", "_fp16", ".fp16", "-q4f16", "_q4f16", ".q4f16", "-q8", "_q8", ".q8", "-bnb4", + "_bnb4", ".bnb4", + ] { + if let Some(base) = stem.strip_suffix(suffix) { + return format!("{base}-{dtype}.onnx"); + } + } + + format!("{stem}-{dtype}.onnx") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_model_path_uses_dtype_suffix_once() { + assert_eq!( + select_local_model_path("/tmp/ecpose.onnx", DType::Fp16), + "/tmp/ecpose-fp16.onnx" + ); + assert_eq!( + select_local_model_path("/tmp/ecpose-fp16.onnx", DType::Fp16), + "/tmp/ecpose-fp16.onnx" + ); + } +} diff --git a/examples/pose-estimation/main.rs b/examples/pose-estimation/main.rs index aa4171e..53d8f16 100644 --- a/examples/pose-estimation/main.rs +++ b/examples/pose-estimation/main.rs @@ -1,12 +1,13 @@ use anyhow::Result; use clap::{Parser, Subcommand}; use usls::{ - models::{CIGPose, DWPose, HRNet, RTMPose, RTMO, YOLO}, + models::{CIGPose, DWPose, ECPose, HRNet, RTMPose, RTMO, YOLO}, Annotator, Config, DataLoader, Model, Scale, Source, Y, }; mod cigpose; mod dwpose; +mod ecpose; mod hrnet; mod rtmo; mod rtmpose; @@ -32,6 +33,7 @@ struct Cli { enum Commands { Cigpose(cigpose::CigposeArgs), Dwpose(dwpose::DwposeArgs), + Ecpose(ecpose::EcposeArgs), Hrnet(hrnet::HrnetArgs), Rtmo(rtmo::RtmoArgs), Rtmpose(rtmpose::RtmposeArgs), @@ -86,6 +88,38 @@ fn main() -> Result<()> { } Ok(()) } + Commands::Ecpose(args) => { + let config = ecpose::config(args)?.commit()?; + let mut model = ECPose::new(config)?; + let annotator = annotator.with_keypoint_style( + usls::KeypointStyle::default() + .with_radius(2) + .with_skeleton((usls::SKELETON_COCO_19, usls::SKELETON_COLOR_COCO_19).into()) + .show_id(false) + .show_confidence(false) + .show_name(false), + ); + let dl = DataLoader::new(&cli.source)? + .with_batch(model.batch() as _) + .with_progress_bar(true) + .stream()?; + + for xs in &dl { + let ys = model.forward(&xs)?; + for (x, y) in xs.iter().zip(ys.iter()) { + if !y.is_empty() { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs/pose-estimation", model.spec()])? + .join(usls::timestamp(None)) + .display() + ))?; + } + } + } + Ok(()) + } Commands::Dwpose(args) => { let yolo_config = yolo_config .with_model_device(args.device) diff --git a/src/models/vision/ecpose/config.rs b/src/models/vision/ecpose/config.rs new file mode 100644 index 0000000..ea007f3 --- /dev/null +++ b/src/models/vision/ecpose/config.rs @@ -0,0 +1,64 @@ +use crate::{ResizeAlg, ResizeFilter, Task, NAMES_COCO_KEYPOINTS_17}; + +const ECPOSE_RELEASE: &str = "https://github.com/wep21/assets/releases/download/ecpose"; + +/// +/// > # ECPose: EdgeCrafter Multi-Person Pose Estimation +/// > +/// > EdgeCrafter pose models exported to ONNX with DETR-style fixed-query outputs. +/// > +/// > # Model Variants +/// > +/// > - **ecpose-s**: Small multi-person COCO-17 pose model +/// > - **ecpose-m**: Medium multi-person COCO-17 pose model +/// > - **ecpose-l**: Large multi-person COCO-17 pose model +/// > - **ecpose-x**: Extra large multi-person COCO-17 pose model +/// > +/// > # Precision / File naming +/// > +/// > Assets are hosted on the `wep21/assets` GitHub releases (tag `ecpose`). +/// > FP32 weights use `ecpose-{scale}.onnx`; non-FP32 weights follow the standard +/// > automatic suffix resolution via [`crate::Config::with_model_dtype`]. +/// +/// Model configuration for `ECPose` +/// +impl crate::Config { + /// Base configuration for ECPose models. + pub fn ecpose() -> Self { + Self::default() + .with_name("ecpose") + .with_task(Task::KeypointsDetection) + .with_model_ixx(0, 0, 1) + .with_model_ixx(0, 1, 3) + .with_model_ixx(0, 2, 640) + .with_model_ixx(0, 3, 640) + .with_model_ixx(1, 1, 2) + .with_resize_mode_type(crate::ResizeModeType::FitExact) + .with_resize_alg(ResizeAlg::Interpolation(ResizeFilter::Bilinear)) + .with_image_mean([0.485, 0.456, 0.406]) + .with_image_std([0.229, 0.224, 0.225]) + .with_class_confs(&[0.5]) + .with_nk(17) + .with_keypoint_names(&NAMES_COCO_KEYPOINTS_17) + } + + /// Small multi-person COCO-17 pose model. + pub fn ecpose_s() -> Self { + Self::ecpose().with_model_file(format!("{ECPOSE_RELEASE}/ecpose-s.onnx")) + } + + /// Medium multi-person COCO-17 pose model. + pub fn ecpose_m() -> Self { + Self::ecpose().with_model_file(format!("{ECPOSE_RELEASE}/ecpose-m.onnx")) + } + + /// Large multi-person COCO-17 pose model. + pub fn ecpose_l() -> Self { + Self::ecpose().with_model_file(format!("{ECPOSE_RELEASE}/ecpose-l.onnx")) + } + + /// Extra large multi-person COCO-17 pose model. + pub fn ecpose_x() -> Self { + Self::ecpose().with_model_file(format!("{ECPOSE_RELEASE}/ecpose-x.onnx")) + } +} diff --git a/src/models/vision/ecpose/impl.rs b/src/models/vision/ecpose/impl.rs new file mode 100644 index 0000000..61cd55c --- /dev/null +++ b/src/models/vision/ecpose/impl.rs @@ -0,0 +1,374 @@ +use aksr::Builder; +use anyhow::{ensure, Result}; +use rayon::prelude::*; + +use crate::{ + inputs, Config, DynConf, Engine, Engines, FromConfig, Image, ImageProcessor, Keypoint, Model, + Module, ResizeModeType, XAny, XView, Xs, X, Y, +}; + +/// ECPose: EdgeCrafter multi-person pose estimation. +/// +/// The ONNX export accepts `images` `[N, 3, 640, 640]` and `orig_target_sizes` +/// `[N, 2]` in `[width, height]` order, then returns fixed-query pose tensors: +/// `scores [N, 60]`, `labels [N, 60]`, and `keypoints [N, 60, 17, 2]`. +#[derive(Debug, Builder)] +pub struct ECPose { + pub height: usize, + pub width: usize, + pub batch: usize, + pub max_queries: usize, + pub num_keypoints: usize, + pub confs: DynConf, + pub names: Vec, + pub processor: ImageProcessor, + pub spec: String, +} + +impl Model for ECPose { + type Input<'a> = &'a [Image]; + + fn batch(&self) -> usize { + self.batch + } + + fn spec(&self) -> &str { + &self.spec + } + + fn build(mut config: Config) -> Result<(Self, Engines)> { + let engine = Engine::from_config(config.take_module(&Module::Model)?)?; + let (batch, height, width) = ( + engine.batch().opt(), + engine.try_height().unwrap_or(&640.into()).opt(), + engine.try_width().unwrap_or(&640.into()).opt(), + ); + let spec = engine.spec().to_owned(); + let num_keypoints = config.inference.num_keypoints.unwrap_or(17); + let names = config.inference.keypoint_names; + let confs = DynConf::new_or(&config.inference.class_confs, 1, 0.5); + let processor = ImageProcessor::from_config(config.image_processor)? + .with_image_width(width as _) + .with_image_height(height as _); + + let model = Self { + height, + width, + batch, + max_queries: 60, + num_keypoints, + confs, + names, + processor, + spec, + }; + + Ok((model, Engines::from(engine))) + } + + fn run(&mut self, engines: &mut Engines, images: Self::Input<'_>) -> Result> { + let xs_images = crate::perf!("ECPose::preprocess", self.processor.process(images)?); + self.validate_image_input(&xs_images, images.len())?; + + let orig_target_sizes = + X::::from_shape_vec_generic([images.len(), 2], self.orig_target_sizes(images))?; + self.validate_orig_target_sizes(&orig_target_sizes, images.len())?; + + let ys = crate::perf!( + "ECPose::inference", + engines.run(&Module::Model, inputs![&xs_images, orig_target_sizes]?)? + ); + crate::perf!("ECPose::postprocess", self.postprocess(&ys)) + } +} + +impl ECPose { + fn orig_target_sizes(&self, images: &[Image]) -> Vec { + images + .iter() + .flat_map(|image| [image.width() as i64, image.height() as i64]) + .collect() + } + + fn validate_image_input(&self, xs_images: &XAny, batch: usize) -> Result<()> { + ensure!( + matches!( + self.processor.resize_mode_type(), + Some(ResizeModeType::FitExact) + ), + "ECPose expects FitExact preprocessing, got {:?}", + self.processor.resize_mode_type() + ); + + let shape = xs_images.shape(); + ensure!( + shape == [batch, 3, self.height, self.width], + "ECPose expects image input shape [{batch}, 3, {}, {}], got {shape:?}", + self.height, + self.width + ); + Ok(()) + } + + fn validate_orig_target_sizes(&self, orig_target_sizes: &X, batch: usize) -> Result<()> { + let shape = orig_target_sizes.dims(); + ensure!( + shape == [batch, 2], + "ECPose expects orig_target_sizes shape [{batch}, 2], got {shape:?}" + ); + Ok(()) + } + + fn postprocess(&self, outputs: &Xs) -> Result> { + ensure!( + outputs.len() == 3, + "ECPose expects exactly 3 outputs (scores, labels, keypoints), got {}", + outputs.len() + ); + + let scores = outputs + .get::(0) + .ok_or_else(|| anyhow::anyhow!("Failed to get ECPose scores output"))?; + let labels = outputs + .get::(1) + .ok_or_else(|| anyhow::anyhow!("Failed to get ECPose labels output"))?; + let keypoints = outputs + .get::(2) + .ok_or_else(|| anyhow::anyhow!("Failed to get ECPose keypoints output"))?; + + self.postprocess_tensors(scores, labels, keypoints) + } + + fn postprocess_tensors( + &self, + scores: XView<'_, f32>, + labels: XView<'_, i64>, + keypoints: XView<'_, f32>, + ) -> Result> { + let (batch, queries, num_keypoints) = + self.validate_output_shapes(&scores, &labels, &keypoints)?; + ensure!( + self.processor.transform_info().len() == batch, + "ECPose expected {batch} image transform records, got {}", + self.processor.transform_info().len() + ); + + let has_names = !self.names.is_empty(); + let ys = (0..batch) + .into_par_iter() + .map(|batch_idx| { + let info = &self.processor.transform_info()[batch_idx]; + let width = info.width_src as f32; + let height = info.height_src as f32; + + let keypointss = (0..queries) + .filter_map(|query_idx| { + let score = scores[[batch_idx, query_idx]]; + let label = labels[[batch_idx, query_idx]]; + if label < 0 { + return None; + } + + let class_id = label as usize; + if score < self.class_conf(class_id) { + return None; + } + + let pose = (0..num_keypoints) + .map(|keypoint_idx| { + let x = keypoints[[batch_idx, query_idx, keypoint_idx, 0]] + .clamp(0.0, width); + let y = keypoints[[batch_idx, query_idx, keypoint_idx, 1]] + .clamp(0.0, height); + let mut keypoint = Keypoint::default() + .with_xy(x, y) + .with_id(keypoint_idx) + .with_confidence(score); + if has_names { + keypoint = keypoint.with_name(&self.names[keypoint_idx]); + } + keypoint + }) + .collect::>(); + + Some(pose) + }) + .collect::>(); + + Y::default().with_keypointss(&keypointss) + }) + .collect(); + + Ok(ys) + } + + fn validate_output_shapes( + &self, + scores: &XView<'_, f32>, + labels: &XView<'_, i64>, + keypoints: &XView<'_, f32>, + ) -> Result<(usize, usize, usize)> { + let scores_shape = scores.dims(); + let labels_shape = labels.dims(); + let keypoints_shape = keypoints.dims(); + + ensure!( + scores_shape.len() == 2, + "ECPose expects scores shape [batch, {}], got {scores_shape:?}", + self.max_queries + ); + ensure!( + labels_shape.len() == 2, + "ECPose expects labels shape [batch, {}], got {labels_shape:?}", + self.max_queries + ); + ensure!( + keypoints_shape.len() == 4, + "ECPose expects keypoints shape [batch, {}, {}, 2], got {keypoints_shape:?}", + self.max_queries, + self.num_keypoints + ); + + let batch = scores_shape[0]; + let queries = scores_shape[1]; + ensure!( + labels_shape == scores_shape, + "ECPose expects labels shape to match scores shape, got scores {scores_shape:?} and labels {labels_shape:?}" + ); + ensure!( + queries == self.max_queries, + "ECPose expects {} pose queries, got {}", + self.max_queries, + queries + ); + ensure!( + keypoints_shape[0] == batch, + "ECPose expects keypoints batch dimension {batch}, got {}", + keypoints_shape[0] + ); + ensure!( + keypoints_shape[1] == queries, + "ECPose expects keypoints query dimension {queries}, got {}", + keypoints_shape[1] + ); + ensure!( + keypoints_shape[2] == self.num_keypoints, + "ECPose expects {} keypoints per pose, got {}", + self.num_keypoints, + keypoints_shape[2] + ); + ensure!( + keypoints_shape[3] == 2, + "ECPose expects 2 coordinates per keypoint, got {}", + keypoints_shape[3] + ); + if has_names_mismatch(&self.names, self.num_keypoints) { + anyhow::bail!( + "ECPose expected {} keypoint names or none, got {}", + self.num_keypoints, + self.names.len() + ); + } + + Ok((batch, queries, self.num_keypoints)) + } + + fn class_conf(&self, class_id: usize) -> f32 { + self.confs + .values() + .get(class_id) + .copied() + .or_else(|| self.confs.values().first().copied()) + .unwrap_or(0.5) + } +} + +fn has_names_mismatch(names: &[String], expected: usize) -> bool { + !names.is_empty() && names.len() != expected +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ImageProcessorConfig, ImageTransformInfo, ResizeModeType}; + + fn test_model() -> ECPose { + let mut processor = ImageProcessor::from_config( + ImageProcessorConfig::default().with_resize_mode_type(ResizeModeType::FitExact), + ) + .expect("processor"); + processor.images_transform_info = vec![ImageTransformInfo::default() + .with_width_src(100) + .with_height_src(80)]; + + ECPose { + height: 640, + width: 640, + batch: 1, + max_queries: 60, + num_keypoints: 17, + confs: DynConf::new_or(&[0.5], 1, 0.5), + names: crate::NAMES_COCO_KEYPOINTS_17 + .iter() + .map(|name| (*name).to_string()) + .collect(), + processor, + spec: "ecpose/test".to_string(), + } + } + + #[test] + fn postprocess_filters_and_clamps_keypoints() { + let model = test_model(); + + let mut scores = vec![0.1f32; 60]; + scores[0] = 0.9; + scores[1] = 0.49; + let scores = X::from_shape_vec(&[1, 60], scores).expect("scores"); + + let labels = X::::from_shape_vec_generic([1, 60], vec![0i64; 60]).expect("labels"); + + let mut keypoints = vec![0.0f32; 60 * 17 * 2]; + for idx in 0..17 { + let base = idx * 2; + keypoints[base] = if idx == 0 { -4.0 } else { 10.0 + idx as f32 }; + keypoints[base + 1] = if idx == 0 { 120.0 } else { 20.0 + idx as f32 }; + } + let keypoints = X::from_shape_vec(&[1, 60, 17, 2], keypoints).expect("synthetic keypoints"); + + let ys = model + .postprocess_tensors(scores.view(), labels.view(), keypoints.view()) + .expect("postprocess"); + + assert_eq!(ys.len(), 1); + assert_eq!(ys[0].keypointss.len(), 1); + assert_eq!(ys[0].hbbs.len(), 0); + + let pose = &ys[0].keypointss[0]; + assert_eq!(pose.len(), 17); + assert_eq!(pose[0].id(), Some(0)); + assert_eq!(pose[0].name(), Some("nose")); + assert_eq!(pose[0].confidence(), Some(0.9)); + assert_eq!(pose[0].xy(), (0.0, 80.0)); + assert_eq!(pose[16].confidence(), Some(0.9)); + } + + #[test] + fn postprocess_rejects_invalid_shapes() { + let model = test_model(); + let scores = X::from_shape_vec(&[1, 60], vec![0.5f32; 60]).expect("scores"); + let labels = X::::from_shape_vec_generic([1, 60], vec![0i64; 60]).expect("labels"); + let keypoints = + X::from_shape_vec(&[1, 59, 17, 2], vec![0.0f32; 59 * 17 * 2]).expect("keypoints"); + + let err = model + .postprocess_tensors(scores.view(), labels.view(), keypoints.view()) + .expect_err("shape validation should fail"); + + assert!( + err.to_string() + .contains("ECPose expects keypoints query dimension 60"), + "unexpected error: {err}" + ); + } +} diff --git a/src/models/vision/ecpose/mod.rs b/src/models/vision/ecpose/mod.rs new file mode 100644 index 0000000..fbd2b75 --- /dev/null +++ b/src/models/vision/ecpose/mod.rs @@ -0,0 +1,4 @@ +mod config; +mod r#impl; + +pub use r#impl::*; diff --git a/src/models/vision/mod.rs b/src/models/vision/mod.rs index caebcec..1644641 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`, `cigpose` +//! - **Pose**: `rtmpose`, `rtmw`, `dwpose`, `rtmo`, `hrnet`, `cigpose`, `ecpose` //! - **Depth**: `depth_anything`, `depth_pro` //! - **Feature**: `dinov2`, `dinov3`, `clip`, `blip`, `ram` //! - **OCR**: `db`, `fast`, `linknet`, `svtr`, `slanet` @@ -47,6 +47,7 @@ mod yoloe_prompt_free; // Pose Estimation mod cigpose; mod dwpose; +mod ecpose; mod hrnet; mod rtmo; mod rtmpose; @@ -108,6 +109,7 @@ pub use dinov2::*; pub use dinov3::*; pub use dwpose::*; pub use ecdetseg::*; +pub use ecpose::*; pub use fast::*; pub use fastsam::*; pub use fastvit::*;