From 68a70e6e7bba2a12351c526c5189efbbfd080b6f Mon Sep 17 00:00:00 2001 From: Jamjamjon <51357717+jamjamjon@users.noreply.github.com> Date: Fri, 9 May 2025 17:17:52 +0800 Subject: [PATCH] Add RF-DETR large model (#87) --- src/inference/hbb.rs | 40 ++-------------------------------- src/inference/instance_meta.rs | 34 +++++++++++++++++++++++++++++ src/inference/keypoint.rs | 40 ++-------------------------------- src/inference/mask.rs | 40 ++-------------------------------- src/inference/obb.rs | 40 ++-------------------------------- src/inference/polygon.rs | 40 ++-------------------------------- src/inference/prob.rs | 39 ++------------------------------- src/models/rfdetr/config.rs | 2 +- 8 files changed, 47 insertions(+), 228 deletions(-) diff --git a/src/inference/hbb.rs b/src/inference/hbb.rs index 63402c0..12bcac3 100644 --- a/src/inference/hbb.rs +++ b/src/inference/hbb.rs @@ -1,6 +1,6 @@ use aksr::Builder; -use crate::{InstanceMeta, Keypoint, Style}; +use crate::{impl_meta_methods, InstanceMeta, Keypoint, Style}; #[derive(Builder, Clone, Default)] pub struct Hbb { @@ -83,44 +83,8 @@ impl From for [f32; 4] { } impl Hbb { - pub fn with_uid(mut self, uid: usize) -> Self { - self.meta = self.meta.with_uid(uid); - self - } + impl_meta_methods!(); - pub fn with_id(mut self, id: usize) -> Self { - self.meta = self.meta.with_id(id); - self - } - - pub fn with_name(mut self, name: &str) -> Self { - self.meta = self.meta.with_name(name); - self - } - - pub fn with_confidence(mut self, confidence: f32) -> Self { - self.meta = self.meta.with_confidence(confidence); - self - } - - pub fn uid(&self) -> usize { - self.meta.uid() - } - - pub fn name(&self) -> Option<&str> { - self.meta.name() - } - - pub fn confidence(&self) -> Option { - self.meta.confidence() - } - - pub fn id(&self) -> Option { - self.meta.id() - } -} - -impl Hbb { pub fn from_xywh(x: f32, y: f32, w: f32, h: f32) -> Self { Self { x, diff --git a/src/inference/instance_meta.rs b/src/inference/instance_meta.rs index d67bd39..ef149eb 100644 --- a/src/inference/instance_meta.rs +++ b/src/inference/instance_meta.rs @@ -77,3 +77,37 @@ impl InstanceMeta { label } } + +#[macro_export] +macro_rules! impl_meta_methods { + () => { + pub fn with_uid(mut self, uid: usize) -> Self { + self.meta = self.meta.with_uid(uid); + self + } + pub fn with_id(mut self, id: usize) -> Self { + self.meta = self.meta.with_id(id); + self + } + pub fn with_name(mut self, name: &str) -> Self { + self.meta = self.meta.with_name(name); + self + } + pub fn with_confidence(mut self, confidence: f32) -> Self { + self.meta = self.meta.with_confidence(confidence); + self + } + pub fn uid(&self) -> usize { + self.meta.uid() + } + pub fn name(&self) -> Option<&str> { + self.meta.name() + } + pub fn confidence(&self) -> Option { + self.meta.confidence() + } + pub fn id(&self) -> Option { + self.meta.id() + } + }; +} diff --git a/src/inference/keypoint.rs b/src/inference/keypoint.rs index 8bc6e8d..3bf25e4 100644 --- a/src/inference/keypoint.rs +++ b/src/inference/keypoint.rs @@ -1,7 +1,7 @@ use aksr::Builder; use std::ops::{Add, Div, Mul, Sub}; -use crate::{InstanceMeta, Style}; +use crate::{impl_meta_methods, InstanceMeta, Style}; /// Represents a keypoint in a 2D space with optional metadata. #[derive(Builder, Default, Clone)] @@ -69,44 +69,8 @@ impl From for [f32; 2] { } impl Keypoint { - pub fn with_uid(mut self, uid: usize) -> Self { - self.meta = self.meta.with_uid(uid); - self - } + impl_meta_methods!(); - pub fn with_id(mut self, id: usize) -> Self { - self.meta = self.meta.with_id(id); - self - } - - pub fn with_name(mut self, name: &str) -> Self { - self.meta = self.meta.with_name(name); - self - } - - pub fn with_confidence(mut self, confidence: f32) -> Self { - self.meta = self.meta.with_confidence(confidence); - self - } - - pub fn uid(&self) -> usize { - self.meta.uid() - } - - pub fn name(&self) -> Option<&str> { - self.meta.name() - } - - pub fn confidence(&self) -> Option { - self.meta.confidence() - } - - pub fn id(&self) -> Option { - self.meta.id() - } -} - -impl Keypoint { pub fn new(x: f32, y: f32) -> Self { Self { x, diff --git a/src/inference/mask.rs b/src/inference/mask.rs index 4515a27..09ce1bf 100644 --- a/src/inference/mask.rs +++ b/src/inference/mask.rs @@ -3,7 +3,7 @@ use anyhow::Result; use image::GrayImage; use rayon::prelude::*; -use crate::{InstanceMeta, Polygon, Style}; +use crate::{impl_meta_methods, InstanceMeta, Polygon, Style}; /// Mask: Gray Image. #[derive(Builder, Default, Clone)] @@ -35,6 +35,7 @@ impl PartialEq for Mask { } impl Mask { + impl_meta_methods!(); pub fn new(u8s: &[u8], width: u32, height: u32) -> Result { let mask: image::ImageBuffer, Vec<_>> = image::ImageBuffer::from_raw(width, height, u8s.to_vec()) @@ -103,40 +104,3 @@ impl Mask { polygons } } - -impl Mask { - pub fn with_uid(mut self, uid: usize) -> Self { - self.meta = self.meta.with_uid(uid); - self - } - pub fn with_id(mut self, id: usize) -> Self { - self.meta = self.meta.with_id(id); - self - } - - pub fn with_name(mut self, name: &str) -> Self { - self.meta = self.meta.with_name(name); - self - } - - pub fn with_confidence(mut self, confidence: f32) -> Self { - self.meta = self.meta.with_confidence(confidence); - self - } - - pub fn uid(&self) -> usize { - self.meta.uid() - } - - pub fn name(&self) -> Option<&str> { - self.meta.name() - } - - pub fn confidence(&self) -> Option { - self.meta.confidence() - } - - pub fn id(&self) -> Option { - self.meta.id() - } -} diff --git a/src/inference/obb.rs b/src/inference/obb.rs index f4a6d39..0f2f781 100644 --- a/src/inference/obb.rs +++ b/src/inference/obb.rs @@ -1,6 +1,6 @@ use aksr::Builder; -use crate::{Hbb, InstanceMeta, Keypoint, Polygon, Style}; +use crate::{impl_meta_methods, Hbb, InstanceMeta, Keypoint, Polygon, Style}; #[derive(Builder, Default, Clone, PartialEq)] pub struct Obb { @@ -48,6 +48,7 @@ impl From for [[f32; 2]; 4] { } impl Obb { + impl_meta_methods!(); /// Build from (cx, cy, width, height, degrees) pub fn from_cxcywhd(cx: f32, cy: f32, w: f32, h: f32, d: f32) -> Self { Self::from_cxcywhr(cx, cy, w, h, d.to_radians()) @@ -142,43 +143,6 @@ impl Obb { } } -impl Obb { - pub fn with_uid(mut self, uid: usize) -> Self { - self.meta = self.meta.with_uid(uid); - self - } - pub fn with_id(mut self, id: usize) -> Self { - self.meta = self.meta.with_id(id); - self - } - - pub fn with_name(mut self, name: &str) -> Self { - self.meta = self.meta.with_name(name); - self - } - - pub fn with_confidence(mut self, confidence: f32) -> Self { - self.meta = self.meta.with_confidence(confidence); - self - } - - pub fn uid(&self) -> usize { - self.meta.uid() - } - - pub fn name(&self) -> Option<&str> { - self.meta.name() - } - - pub fn confidence(&self) -> Option { - self.meta.confidence() - } - - pub fn id(&self) -> Option { - self.meta.id() - } -} - #[cfg(test)] mod tests_mbr { // use crate::Nms; diff --git a/src/inference/polygon.rs b/src/inference/polygon.rs index dd3e059..a270f68 100644 --- a/src/inference/polygon.rs +++ b/src/inference/polygon.rs @@ -4,7 +4,7 @@ use geo::{ Point, Simplify, }; -use crate::{Hbb, InstanceMeta, Mask, Obb, Style}; +use crate::{impl_meta_methods, Hbb, InstanceMeta, Mask, Obb, Style}; /// Polygon. #[derive(Builder, Clone)] @@ -43,6 +43,7 @@ impl PartialEq for Polygon { } impl Polygon { + impl_meta_methods!(); pub fn with_points_imageproc(mut self, points: &[imageproc::point::Point]) -> Self { // exterior let v = points @@ -245,40 +246,3 @@ impl Polygon { xs.retain(|point| seen.insert((point.x() as i32, point.y() as i32))); } } - -impl Polygon { - pub fn with_uid(mut self, uid: usize) -> Self { - self.meta = self.meta.with_uid(uid); - self - } - pub fn with_id(mut self, id: usize) -> Self { - self.meta = self.meta.with_id(id); - self - } - - pub fn with_name(mut self, name: &str) -> Self { - self.meta = self.meta.with_name(name); - self - } - - pub fn with_confidence(mut self, confidence: f32) -> Self { - self.meta = self.meta.with_confidence(confidence); - self - } - - pub fn uid(&self) -> usize { - self.meta.uid() - } - - pub fn name(&self) -> Option<&str> { - self.meta.name() - } - - pub fn confidence(&self) -> Option { - self.meta.confidence() - } - - pub fn id(&self) -> Option { - self.meta.id() - } -} diff --git a/src/inference/prob.rs b/src/inference/prob.rs index 70a93b3..107859c 100644 --- a/src/inference/prob.rs +++ b/src/inference/prob.rs @@ -1,6 +1,6 @@ use aksr::Builder; -use crate::{InstanceMeta, Style}; +use crate::{impl_meta_methods, InstanceMeta, Style}; #[derive(Builder, Clone, PartialEq, Default, Debug)] pub struct Prob { @@ -12,6 +12,7 @@ pub struct Prob { // pub struct Probs(#[args(aka = "probs")] Vec); impl Prob { + impl_meta_methods!(); pub fn new_probs(probs: &[f32], names: Option<&[&str]>, k: usize) -> Vec { let mut pairs: Vec<(usize, f32)> = probs .iter() @@ -39,40 +40,4 @@ impl Prob { }) .collect() } - - pub fn with_uid(mut self, uid: usize) -> Self { - self.meta = self.meta.with_uid(uid); - self - } - - pub fn with_id(mut self, id: usize) -> Self { - self.meta = self.meta.with_id(id); - self - } - - pub fn with_name(mut self, name: &str) -> Self { - self.meta = self.meta.with_name(name); - self - } - - pub fn with_confidence(mut self, confidence: f32) -> Self { - self.meta = self.meta.with_confidence(confidence); - self - } - - pub fn uid(&self) -> usize { - self.meta.uid() - } - - pub fn name(&self) -> Option<&str> { - self.meta.name() - } - - pub fn confidence(&self) -> Option { - self.meta.confidence() - } - - pub fn id(&self) -> Option { - self.meta.id() - } } diff --git a/src/models/rfdetr/config.rs b/src/models/rfdetr/config.rs index 6edccd0..53a861a 100644 --- a/src/models/rfdetr/config.rs +++ b/src/models/rfdetr/config.rs @@ -19,7 +19,7 @@ impl crate::Options { pub fn rfdetr_base() -> Self { Self::rfdetr().with_model_file("base.onnx") } - + pub fn rfdetr_large() -> Self { Self::rfdetr().with_model_file("large.onnx") }