diff --git a/examples/pose-estimation/hrnet.rs b/examples/pose-estimation/hrnet.rs new file mode 100644 index 0000000..e99465c --- /dev/null +++ b/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, + + /// 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 { + 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) +} diff --git a/examples/pose-estimation/main.rs b/examples/pose-estimation/main.rs index f5a86a2..538a3e3 100644 --- a/examples/pose-estimation/main.rs +++ b/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_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) diff --git a/src/models/vision/hrnet/config.rs b/src/models/vision/hrnet/config.rs new file mode 100644 index 0000000..ed89373 --- /dev/null +++ b/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", + ) + } +} diff --git a/src/models/vision/hrnet/impl.rs b/src/models/vision/hrnet/impl.rs new file mode 100644 index 0000000..6045f40 --- /dev/null +++ b/src/models/vision/hrnet/impl.rs @@ -0,0 +1,379 @@ +use aksr::Builder; +use anyhow::Result; +use ndarray::s; +use rayon::prelude::*; +use std::f32::consts::PI; + +use crate::{ + Config, DynConf, Engine, Engines, FromConfig, Hbb, Image, ImageProcessor, Keypoint, Model, + Module, XAny, Xs, X, Y, +}; + +struct CentersAndScales { + pub centers: Vec, + pub scales: Vec, +} + +/// HRNet: top-down heatmap-based human pose estimation +/// +/// Supports COCO 17-keypoint body and COCO-WholeBody 133-keypoint estimation with the +/// `w32` / `w48` backbones. The model takes a person crop and outputs per-keypoint +/// heatmaps of shape `[batch, num_keypoints, h, w]` (1/4 of the input resolution). +#[derive(Builder, Debug)] +pub struct HRNet { + pub height: usize, + pub width: usize, + pub batch: usize, + pub spec: String, + pub processor: ImageProcessor, + pub nk: usize, + pub kconfs: DynConf, + pub names: Vec, +} + +impl Model for HRNet { + 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 spec = engine.spec().to_string(); + let (batch, height, width) = ( + engine.batch().opt(), + engine.try_height().unwrap_or(&256.into()).opt(), + engine.try_width().unwrap_or(&192.into()).opt(), + ); + let nk = config.inference.num_keypoints.unwrap_or(17); + let kconfs = DynConf::new_or_default(&config.inference.keypoint_confs, nk); + let names = config.inference.keypoint_names; + let processor = ImageProcessor::from_config(config.image_processor)? + .with_image_width(width as _) + .with_image_height(height as _); + + let model = Self { + height, + width, + batch, + spec, + processor, + nk, + kconfs, + names, + }; + + let engines = Engines::from(engine); + Ok((model, engines)) + } + + fn run(&mut self, engines: &mut Engines, input: Self::Input<'_>) -> Result> { + let images = input; + let (xs, centers_and_scales) = crate::perf!("HRNet::preprocess", self.preprocess(images)?); + let ys = crate::perf!("HRNet::inference", engines.run(&Module::Model, &xs)?); + let y = crate::perf!("HRNet::postprocess", { + self.postprocess(&ys, centers_and_scales)? + }); + + Ok(y) + } +} + +impl HRNet { + fn preprocess(&mut self, images: &[Image]) -> Result<(XAny, CentersAndScales)> { + let model_input_size = (self.width as i32, self.height as i32); + let results: Vec<(Image, Keypoint, Keypoint)> = images + .par_iter() + .map(|img| { + let hbb = Hbb::from_xyxy(0.0, 0.0, img.width() as f32, img.height() as f32); + let (center, scale) = Self::hbb2cs(&hbb, 1.25); + let (resized_img, scale) = + Self::top_down_affine(model_input_size, &scale, ¢er, img)?; + Ok((resized_img, center, scale)) + }) + .collect::>>()?; + let (processed_images, centers, scales): (Vec<_>, Vec<_>, Vec<_>) = + results.into_iter().fold( + (Vec::new(), Vec::new(), Vec::new()), + |(mut imgs, mut ctrs, mut scls), (img, ctr, scl)| { + imgs.push(img); + ctrs.push(ctr); + scls.push(scl); + (imgs, ctrs, scls) + }, + ); + let x = self.processor.process(&processed_images)?; + self.batch = processed_images.len(); + + Ok((x, CentersAndScales { centers, scales })) + } + + fn postprocess( + &mut self, + outputs: &Xs, + centers_and_scales: CentersAndScales, + ) -> Result> { + let x0 = outputs + .get::(0) + .ok_or_else(|| anyhow::anyhow!("Failed to get output 0"))?; + let heatmaps = X::from(x0); + + // [batch, num_keypoints, hm_h, hm_w] + let shape = heatmaps.shape(); + anyhow::ensure!( + shape.len() == 4, + "HRNet expects a 4D heatmap output [batch, nk, h, w], got shape {shape:?}", + ); + let total_crops = shape[0]; + self.nk = shape[1]; + let (hm_h, hm_w) = (shape[2], shape[3]); + let has_names = !self.names.is_empty(); + + let results: Vec = (0..total_crops) + .into_par_iter() + .map(|batch_idx| { + let center = ¢ers_and_scales.centers[batch_idx]; + let scale = ¢ers_and_scales.scales[batch_idx]; + + let x_factor = scale.x() / hm_w as f32; + let y_factor = scale.y() / hm_h as f32; + let x_offset = center.x() - scale.x() * 0.5; + let y_offset = center.y() - scale.y() * 0.5; + + let keypoints: Vec = (0..self.nk) + .map(|kpt_idx| { + let hm = heatmaps.slice(s![batch_idx, kpt_idx, .., ..]); + let (py, px, confidence) = Self::heatmap_argmax(&hm); + + if confidence > self.kconfs[kpt_idx] { + // Sub-pixel refinement: shift 0.25 toward the larger neighbor. + let (mut fx, mut fy) = (px as f32, py as f32); + if px > 0 && px < hm_w - 1 && py > 0 && py < hm_h - 1 { + let dx = hm[[py, px + 1]] - hm[[py, px - 1]]; + let dy = hm[[py + 1, px]] - hm[[py - 1, px]]; + fx += dx.signum() * 0.25; + fy += dy.signum() * 0.25; + } + + let x = fx * x_factor + x_offset; + let y = fy * y_factor + y_offset; + + let mut kpt = Keypoint::from((x, y)) + .with_confidence(confidence) + .with_id(kpt_idx); + + if has_names { + kpt = kpt.with_name(&self.names[kpt_idx]); + } + kpt + } else { + Keypoint::default() + } + }) + .collect(); + + Y::default().with_keypointss(&[keypoints]) + }) + .collect(); + + Ok(results) + } + + fn heatmap_argmax(hm: &ndarray::ArrayView2) -> (usize, usize, f32) { + let mut max_val = f32::MIN; + let (mut max_y, mut max_x) = (0usize, 0usize); + for ((y, x), &val) in hm.indexed_iter() { + if val > max_val { + max_val = val; + max_y = y; + max_x = x; + } + } + + (max_y, max_x, max_val) + } + + fn hbb2cs(hbb: &Hbb, padding: f32) -> (Keypoint, Keypoint) { + let (x1, y1, x2, y2) = hbb.xyxy(); + ( + ((x1 + x2) * 0.5, (y1 + y2) * 0.5).into(), + ((x2 - x1) * padding, (y2 - y1) * padding).into(), + ) + } + + fn get_warp_matrix( + center: &Keypoint, + scale: &Keypoint, + rot: f32, + output_size: (i32, i32), + shift: (f32, f32), + inv: bool, + ) -> Vec { + fn get_3rd_point(a: &Keypoint, b: &Keypoint) -> Keypoint { + let direction = Keypoint::new(a.x() - b.x(), a.y() - b.y()); + (b.x() - direction.y(), b.y() + direction.x()).into() + } + + let shift = Keypoint::new(shift.0, shift.1); + let src_w = scale.x(); + let dst_w = output_size.0 as f32; + let dst_h = output_size.1 as f32; + let rot_rad = rot * PI / 180.0; + let src_dir = Keypoint::new(0.0, src_w * -0.5).rotate(rot_rad); + let dst_dir = Keypoint::new(0.0, dst_w * -0.5); + let src_0 = Keypoint::new( + center.x() + scale.x() * shift.x(), + center.y() + scale.y() * shift.y(), + ); + let src_1 = Keypoint::new( + center.x() + src_dir.x() + scale.x() * shift.x(), + center.y() + src_dir.y() + scale.y() * shift.y(), + ); + let src_2 = get_3rd_point(&src_0, &src_1); + let dst_0 = Keypoint::new(dst_w * 0.5, dst_h * 0.5); + let dst_1 = Keypoint::new(dst_w * 0.5 + dst_dir.x(), dst_h * 0.5 + dst_dir.y()); + let dst_2 = get_3rd_point(&dst_0, &dst_1); + let (src_points, dst_points) = if inv { + (vec![&dst_0, &dst_1, &dst_2], vec![&src_0, &src_1, &src_2]) + } else { + (vec![&src_0, &src_1, &src_2], vec![&dst_0, &dst_1, &dst_2]) + }; + + let x1 = src_points[0].x(); + let y1 = src_points[0].y(); + let x2 = src_points[1].x(); + let y2 = src_points[1].y(); + let x3 = src_points[2].x(); + let y3 = src_points[2].y(); + + let u1 = dst_points[0].x(); + let v1 = dst_points[0].y(); + let u2 = dst_points[1].x(); + let v2 = dst_points[1].y(); + let u3 = dst_points[2].x(); + let v3 = dst_points[2].y(); + + let det = x1 * (y2 - y3) + x2 * (y3 - y1) + x3 * (y1 - y2); + + if det.abs() < 1e-6 { + return vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0]; + } + + let m00 = (u1 * (y2 - y3) + u2 * (y3 - y1) + u3 * (y1 - y2)) / det; + let m01 = (x1 * (u3 - u2) + x2 * (u1 - u3) + x3 * (u2 - u1)) / det; + let m02 = + (x1 * (y2 * u3 - y3 * u2) + x2 * (y3 * u1 - y1 * u3) + x3 * (y1 * u2 - y2 * u1)) / det; + let m10 = (v1 * (y2 - y3) + v2 * (y3 - y1) + v3 * (y1 - y2)) / det; + let m11 = (x1 * (v2 - v3) + x2 * (v3 - v1) + x3 * (v1 - v2)) / det; + let m12 = + (x1 * (y2 * v3 - y3 * v2) + x2 * (y3 * v1 - y1 * v3) + x3 * (y1 * v2 - y2 * v1)) / det; + + vec![m00, m01, m02, m10, m11, m12] + } + + fn warp_affine(img: &Image, warp_mat: &[f32], output_size: (i32, i32)) -> Result { + let (width, height) = output_size; + let img_w = img.width(); + let img_h = img.height(); + + let m00 = warp_mat[0]; + let m01 = warp_mat[1]; + let m02 = warp_mat[2]; + let m10 = warp_mat[3]; + let m11 = warp_mat[4]; + let m12 = warp_mat[5]; + + let det = m00 * m11 - m01 * m10; + if det.abs() < 1e-6 { + return Image::from_u8s( + &vec![0u8; (height * width * 3) as usize], + width as u32, + height as u32, + ); + } + + let inv_det = 1.0 / det; + let inv_m00 = m11 * inv_det; + let inv_m01 = -m01 * inv_det; + let inv_m10 = -m10 * inv_det; + let inv_m11 = m00 * inv_det; + let inv_m02 = (m01 * m12 - m11 * m02) * inv_det; + let inv_m12 = (m10 * m02 - m00 * m12) * inv_det; + let img_data = img.as_raw(); + + let mut result_data = vec![0u8; (height * width * 3) as usize]; + result_data + .par_chunks_exact_mut((width * 3) as usize) + .enumerate() + .for_each(|(y, row)| { + let dst_y = y as f32; + for x in 0..width { + let dst_x = x as f32; + let src_x = inv_m00 * dst_x + inv_m01 * dst_y + inv_m02; + let src_y = inv_m10 * dst_x + inv_m11 * dst_y + inv_m12; + + let x0 = src_x.floor() as i32; + let y0 = src_y.floor() as i32; + let x1 = x0 + 1; + let y1 = y0 + 1; + + if x0 >= 0 && x1 < img_w as i32 && y0 >= 0 && y1 < img_h as i32 { + let dx = src_x - x0 as f32; + let dy = src_y - y0 as f32; + let w00 = (1.0 - dx) * (1.0 - dy); + let w01 = dx * (1.0 - dy); + let w10 = (1.0 - dx) * dy; + let w11 = dx * dy; + + let dst_idx = (x * 3) as usize; + let y0_offset = (y0 as u32 * img_w) as usize * 3; + let y1_offset = (y1 as u32 * img_w) as usize * 3; + let x0_offset = x0 as usize * 3; + let x1_offset = x1 as usize * 3; + + let src_idx0 = y0_offset + x0_offset; + let src_idx1 = y0_offset + x1_offset; + let src_idx2 = y1_offset + x0_offset; + let src_idx3 = y1_offset + x1_offset; + + for c in 0..3 { + row[dst_idx + c] = (img_data[src_idx0 + c] as f32 * w00 + + img_data[src_idx1 + c] as f32 * w01 + + img_data[src_idx2 + c] as f32 * w10 + + img_data[src_idx3 + c] as f32 * w11) + .round() as u8; + } + } + } + }); + + Image::from_u8s(&result_data, width as u32, height as u32) + } + + fn top_down_affine( + input_size: (i32, i32), + scale: &Keypoint, + center: &Keypoint, + img: &Image, + ) -> Result<(Image, Keypoint)> { + let (w, h) = input_size; + let aspect_ratio = w as f32 / h as f32; + let b_w = scale.x(); + let b_h = scale.y(); + let scale = if b_w > b_h * aspect_ratio { + Keypoint::new(b_w, b_w / aspect_ratio) + } else { + Keypoint::new(b_h * aspect_ratio, b_h) + }; + let rot = 0.0; + let warp_mat = Self::get_warp_matrix(center, &scale, rot, (w, h), (0.0, 0.0), false); + let img = Self::warp_affine(img, &warp_mat, (w, h))?; + + Ok((img, scale)) + } +} diff --git a/src/models/vision/hrnet/mod.rs b/src/models/vision/hrnet/mod.rs new file mode 100644 index 0000000..fbd2b75 --- /dev/null +++ b/src/models/vision/hrnet/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 0df9088..82808e2 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` +//! - **Pose**: `rtmpose`, `rtmw`, `dwpose`, `rtmo`, `hrnet` //! - **Depth**: `depth_anything`, `depth_pro` //! - **Feature**: `dinov2`, `dinov3`, `clip`, `blip`, `ram` //! - **OCR**: `db`, `fast`, `linknet`, `svtr`, `slanet` @@ -45,6 +45,7 @@ mod yoloe_prompt_free; // Pose Estimation mod dwpose; +mod hrnet; mod rtmo; mod rtmpose; mod rtmw; @@ -106,6 +107,7 @@ pub use dwpose::*; pub use fast::*; pub use fastsam::*; pub use fastvit::*; +pub use hrnet::*; pub use linknet::*; pub use mediapipe_segmenter::*; pub use mobile_gaze::*; diff --git a/src/ort/config.rs b/src/ort/config.rs index 9814d9d..dbb52b7 100644 --- a/src/ort/config.rs +++ b/src/ort/config.rs @@ -60,7 +60,7 @@ impl ORTConfig { // Remote match Hub::is_valid_github_release_url(&self.file) { - Some((owner, repo, tag, _file_name)) => { + Some((owner, repo, tag, file_name)) => { // Explicit GitHub release URL detected tracing::debug!( "Explicit GitHub URL detected: {}/{} (tag: {})", @@ -68,9 +68,72 @@ impl ORTConfig { repo, tag ); - let stem = try_fetch_file_stem(&self.file)?; - self.spec = format!("{name}/{owner}-{repo}-{tag}-{stem}"); - self.file = Hub::default().try_fetch(&self.file)?; + + // Build candidate URLs based on dtype, matching the logic + // used for bare filenames in the None branch below. + let candidates: Vec = match self.dtype { + _d @ (DType::Auto | DType::Fp32) => { + vec![self.file.clone()] + } + dtype => { + let mut base = file_name.clone(); + let suffix = base.split_off(base.len() - 5); // 5 -> ".onnx" + ['-', '_', '.'] + .iter() + .map(|delim| { + format!( + "https://github.com/{owner}/{repo}/releases/download/{tag}/{base}{delim}{dtype}{suffix}" + ) + }) + .collect() + } + }; + tracing::debug!( + "Generated {} candidate URLs for resolution: {:?}", + candidates.len(), + candidates + ); + + let mut hub = Hub::default(); + let mut fetch_success = false; + for candidate in &candidates { + // Phase 1: check cache + if let Some(cached_path) = hub.cached(candidate) { + self.file = cached_path; + let stem = try_fetch_file_stem(candidate)?; + self.spec = format!("{name}/{owner}-{repo}-{tag}-{stem}"); + tracing::debug!("Cache hit: {} -> {}", candidate, &self.file); + fetch_success = true; + break; + } + // Phase 2: remote fetch + match hub.try_fetch(candidate) { + Ok(f) => { + self.file = f; + let stem = try_fetch_file_stem(candidate)?; + self.spec = format!("{name}/{owner}-{repo}-{tag}-{stem}"); + tracing::debug!( + "Successfully resolved candidate '{}' to spec: {}", + candidate, + &self.spec + ); + fetch_success = true; + break; + } + Err(err) => { + tracing::warn!("Failed to download candidate '{candidate}': {err}"); + } + } + } + + if !fetch_success { + anyhow::bail!( + "Failed to fetch ONNX model file from URL. \ + None of the generated candidates could be resolved. \ + Please verify the model file path: {:?}", + self.file + ); + } } None => { // Not an explicit GitHub URL — could be a HuggingFace Hub path or