From dc43fd319a480c6af05e66e5ed5c5bc023251009 Mon Sep 17 00:00:00 2001 From: uttarayan21 Date: Thu, 17 Jul 2025 16:18:29 +0530 Subject: [PATCH] feat: added some post processing for retinaface --- Cargo.lock | 27 ++ Cargo.toml | 26 +- bounding-box/Cargo.toml | 7 + bounding-box/src/draw.rs | 39 +++ bounding-box/src/lib.rs | 261 ++++++++++++++++-- bounding-box/src/nms.rs | 68 +++++ .../rust-analyzer/metadata/sysroot/Cargo.lock | 116 ++------ flake.nix | 3 +- src/facedet.rs | 38 ++- 9 files changed, 457 insertions(+), 128 deletions(-) create mode 100644 bounding-box/src/draw.rs create mode 100644 bounding-box/src/nms.rs diff --git a/Cargo.lock b/Cargo.lock index 8912810..6b3b227 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -246,6 +246,17 @@ version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6099cdc01846bc367c4e7dd630dc5966dccf36b652fae7a74e17b640411a91b2" +[[package]] +name = "bounding-box" +version = "0.1.0" +dependencies = [ + "color", + "itertools 0.14.0", + "nalgebra", + "ndarray 0.16.1", + "num", +] + [[package]] name = "built" version = "0.7.7" @@ -386,6 +397,12 @@ dependencies = [ "cc", ] +[[package]] +name = "color" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ae467d04a8a8aea5d9a49018a6ade2e4221d92968e8ce55a48c0b1164e5f698" + [[package]] name = "color_quant" version = "1.1.0" @@ -478,6 +495,7 @@ dependencies = [ name = "detector" version = "0.1.0" dependencies = [ + "bounding-box", "clap", "clap_complete", "error-stack", @@ -961,6 +979,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" diff --git a/Cargo.toml b/Cargo.toml index 222999a..fadc556 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["ndarray-image", "ndarray-resize", "."] +members = ["ndarray-image", "ndarray-resize", ".", "bounding-box"] [workspace.package] version = "0.1.0" @@ -8,6 +8,16 @@ edition = "2024" [workspace.dependencies] ndarray-image = { path = "ndarray-image" } ndarray-resize = { path = "ndarray-resize" } +mnn = { path = "/home/servius/Projects/mnn-rs", version = "0.2.0", features = [ + "opencl", + "tracing", +] } +mnn-bridge = { path = "/home/servius/Projects/mnn-rs/mnn-bridge", version = "0.1.0", features = [ + "ndarray", +] } +mnn-sync = { path = "/home/servius/Projects/mnn-rs/mnn-sync", version = "0.1.0", features = [ + "tracing", +] } [package] name = "detector" @@ -22,16 +32,6 @@ error-stack = "0.5" fast_image_resize = "5.2.0" image = "0.25.6" linfa = "0.7.1" -mnn = { path = "/Users/fs0c131y/Projects/aftershoot/mnn-rs", version = "0.2.0", features = [ - "metal", - "tracing", -] } -mnn-bridge = { path = "/Users/fs0c131y/Projects/aftershoot/mnn-rs/mnn-bridge", version = "0.1.0", features = [ - "ndarray", -] } -mnn-sync = { path = "/Users/fs0c131y/Projects/aftershoot/mnn-rs/mnn-sync", version = "0.1.0", features = [ - "tracing", -] } nalgebra = "0.33.2" ndarray = "0.16.1" ndarray-image = { workspace = true } @@ -42,6 +42,10 @@ thiserror = "2.0" tokio = "1.43.1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } +mnn = { workspace = true } +mnn-bridge = { workspace = true } +mnn-sync = { workspace = true } +bounding-box = { version = "0.1.0", path = "bounding-box" } [profile.release] debug = true diff --git a/bounding-box/Cargo.toml b/bounding-box/Cargo.toml index 8e6c7c7..9d875ba 100644 --- a/bounding-box/Cargo.toml +++ b/bounding-box/Cargo.toml @@ -4,5 +4,12 @@ version = "0.1.0" edition = "2024" [dependencies] +color = "0.3.1" +itertools = "0.14.0" nalgebra = "0.33.2" +ndarray = { version = "0.16.1", optional = true } num = "0.4.3" + +[features] +ndarray = ["dep:ndarray"] +default = ["ndarray"] diff --git a/bounding-box/src/draw.rs b/bounding-box/src/draw.rs new file mode 100644 index 0000000..ada3edd --- /dev/null +++ b/bounding-box/src/draw.rs @@ -0,0 +1,39 @@ +use crate::*; +pub use color::Rgba8; + +use ndarray::{Array3, ArrayViewMut3}; + +pub trait Draw { + fn draw(&mut self, item: T, color: color::Rgba8, thickness: usize); +} + +// impl> Draw for Array3 { +// fn draw(&self, item: T, color: color::Rgba8, thickness: usize) { +// item.draw(&self, color, thickness); +// } +// } + +pub trait Drawable { + fn draw(&self, canvas: &mut Canvas, color: color::Rgba8, thickness: T); +} + +/// Implementing Drawable for Aabb2 with Array3 as the canvas type +/// Assuming Array3 is a 3D array representing an image with RGB/RGBA channels +impl Drawable, T> for Aabb2 +where + T: Num + core::ops::SubAssign + core::ops::AddAssign + core::ops::DivAssign, + T: PartialOrd, +{ + fn draw(&self, canvas: &mut ArrayViewMut3, color: color::Rgba8, thickness: T) { + use itertools::Itertools; + let (height, width, channels) = canvas.dim(); + self.corners() + .iter() + .zip(self.padding(thickness).corners()) + .tuple_windows() + .for_each(|((a, b), (c, d))| { + let bbox = Aabb2::from_vertices([*a, b, *c, d]); + todo!(); + }); + } +} diff --git a/bounding-box/src/lib.rs b/bounding-box/src/lib.rs index c77d163..281ff62 100644 --- a/bounding-box/src/lib.rs +++ b/bounding-box/src/lib.rs @@ -1,26 +1,53 @@ -use nalgebra::{Point, Point2, SVector}; +pub mod draw; +pub mod nms; + +use nalgebra::{Point, Point2, Point3, SVector}; pub trait Num: num::Num + Copy + core::fmt::Debug + 'static {} impl Num for T {} +/// An axis aligned bounding box in `D` dimensions, defined by the minimum vertex and a size vector. #[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct BoundingBox { +pub struct AxisAlignedBoundingBox { /// The point of the bounding box closest to the origin point: Point, /// The size of the bounding box in each dimension size: SVector, } -impl BoundingBox { +pub type Aabb = AxisAlignedBoundingBox; +pub type Aabb2 = AxisAlignedBoundingBox; +pub type Aabb3 = AxisAlignedBoundingBox; + +impl AxisAlignedBoundingBox { pub fn new(point: Point, size: SVector) -> Self { Self { point, size } } - pub fn point(&self) -> &Point { - &self.point + pub fn from_min_max_vertices(point1: Point, point2: Point) -> Self + where + T: core::ops::SubAssign, + { + let size = point2 - point1; + Self::new(point1, SVector::from(size)) } - pub fn size(&self) -> &SVector { - &self.size + pub fn from_vertices(points: [Point; 4]) -> Option + where + T: core::ops::SubAssign, + T: PartialOrd, + { + // find the closest and farthest points from the origin + let min = points + .iter() + .reduce(|acc, p| (acc > p).then_some(p).unwrap_or(acc))?; + let max = points + .iter() + .reduce(|acc, p| (acc < p).then_some(p).unwrap_or(acc))?; + Some(Self::from_min_max_vertices(*min, *max)) + } + + pub fn size(&self) -> SVector { + self.size } pub fn center(&self) -> Point @@ -55,23 +82,109 @@ impl BoundingBox { self.point += translation; } - pub fn contains(&self, point: &Point) -> bool + pub fn min_vertex(&self) -> Point + where + T: core::ops::SubAssign, + { + self.point + } + + pub fn max_vertex(&self) -> Point + where + T: core::ops::AddAssign, + { + self.point + self.size + } + + pub fn contains_point(&self, point: &Point) -> bool where T: core::ops::AddAssign, T: core::ops::SubAssign, T: PartialOrd, { - let p1 = self.point.coords; - let p2 = self.point.coords + self.size; + let min = self.min_vertex(); + let max = self.max_vertex(); - point.coords > p1 && point.coords < p2 + *point > min && *point < max + } + + pub fn scale(self, vector: SVector) -> Self + where + T: core::ops::MulAssign, + T: core::ops::DivAssign, + T: core::ops::SubAssign, + { + let two = T::one() + T::one(); + let new_size = self.size.component_mul(&vector); + let new_point = self.point.coords - new_size / two; + Self { + point: Point::from(new_point), + size: new_size, + } + } + + pub fn contains_bbox(&self, other: &Self) -> bool + where + T: core::ops::AddAssign, + T: core::ops::SubAssign, + T: PartialOrd, + { + let self_min = self.min_vertex(); + let self_max = self.max_vertex(); + let other_min = other.min_vertex(); + let other_max = other.max_vertex(); + + other_min >= self_min && other_max <= self_max + } + + pub fn union(&self, other: &Self) -> Self + where + T: core::ops::AddAssign, + T: core::ops::SubAssign, + T: PartialOrd, + { + let self_min = self.min_vertex(); + let self_max = self.max_vertex(); + let other_min = other.min_vertex(); + let other_max = other.max_vertex(); + let max_of_min = (self_min.coords < other_min.coords) + .then_some(self_min.coords) + .unwrap_or(other_min.coords); + let min_of_max = (self_max.coords > other_max.coords) + .then_some(self_max.coords) + .unwrap_or(other_max.coords); + Self::from_min_max_vertices(Point::from(max_of_min), Point::from(min_of_max)) + } + + pub fn intersection(&self, other: &Self) -> Option + where + T: core::ops::AddAssign, + T: core::ops::SubAssign, + T: PartialOrd, + { + let self_min = self.min_vertex(); + let self_max = self.max_vertex(); + let other_min = other.min_vertex(); + let other_max = other.max_vertex(); + + if self_max < other_min || other_max < self_min { + return None; // No intersection + } + + let min = (self_min.coords > other_min.coords) + .then_some(self_min.coords) + .unwrap_or(other_min.coords); + let max = (self_max.coords < other_max.coords) + .then_some(self_max.coords) + .unwrap_or(other_max.coords); + Some(Self::from_min_max_vertices( + Point::from(min), + Point::from(max), + )) } } -pub type BoundingBox2D = BoundingBox; -pub type BoundingBox3D = BoundingBox; - -impl BoundingBox2D { +impl Aabb2 { pub fn new_2d(point1: Point2, point2: Point2) -> Self where T: core::ops::SubAssign, @@ -110,6 +223,41 @@ impl BoundingBox2D { { [self.x1y1(), self.x2y1(), self.x2y2(), self.x1y2()] } + + pub fn area(&self) -> T + where + T: core::ops::Mul, + { + self.size.x * self.size.y + } + + pub fn iou(&self, other: &Self) -> Option + where + T: core::ops::AddAssign, + T: core::ops::SubAssign, + T: PartialOrd, + { + let intersection = self.intersection(other)?; + let union = self.union(other); + Some(intersection.area() / union.area()) + } +} + +impl Aabb3 { + pub fn new_3d(point1: Point3, point2: Point3) -> Self + where + T: core::ops::SubAssign, + { + let size = point2.coords - point1.coords; + Self::new(point1, SVector::from(size)) + } + + pub fn volume(&self) -> T + where + T: core::ops::Mul, + { + self.size.x * self.size.y * self.size.z + } } #[test] @@ -118,10 +266,10 @@ fn test_bbox_new() { let point1 = Point2::new(1.0, 2.0); let point2 = Point2::new(4.0, 6.0); - let bbox = BoundingBox::new_2d(point1, point2); + let bbox = AxisAlignedBoundingBox::new_2d(point1, point2); - assert_eq!(bbox.point(), &point1); - assert_eq!(bbox.size(), &Vector2::new(3.0, 4.0)); + assert_eq!(bbox.min_vertex(), point1); + assert_eq!(bbox.size(), Vector2::new(3.0, 4.0)); assert_eq!(bbox.center(), Point2::new(2.5, 4.0)); } @@ -131,10 +279,10 @@ fn test_bounding_box_center_2d() { let point = Point2::new(1.0, 2.0); let size = Vector2::new(3.0, 4.0); - let bbox = BoundingBox::new(point, size); + let bbox = AxisAlignedBoundingBox::new(point, size); - assert_eq!(bbox.point(), &point); - assert_eq!(bbox.size(), &size); + assert_eq!(bbox.min_vertex(), point); + assert_eq!(bbox.size(), size); assert_eq!(bbox.center(), Point2::new(2.5, 4.0)); } @@ -144,10 +292,10 @@ fn test_bounding_box_center_3d() { let point = Point3::new(1.0, 2.0, 3.0); let size = Vector3::new(4.0, 5.0, 6.0); - let bbox = BoundingBox::new(point, size); + let bbox = AxisAlignedBoundingBox::new(point, size); - assert_eq!(bbox.point(), &point); - assert_eq!(bbox.size(), &size); + assert_eq!(bbox.min_vertex(), point); + assert_eq!(bbox.size(), size); assert_eq!(bbox.center(), Point3::new(3.0, 4.5, 6.0)); } @@ -157,9 +305,68 @@ fn test_bounding_box_padding_2d() { let point = Point2::new(1.0, 2.0); let size = Vector2::new(3.0, 4.0); - let bbox = BoundingBox::new(point, size); + let bbox = AxisAlignedBoundingBox::new(point, size); let padded_bbox = bbox.padding(1.0); - assert_eq!(padded_bbox.point(), &Point2::new(0.5, 1.5)); - assert_eq!(padded_bbox.size(), &Vector2::new(4.0, 5.0)); + assert_eq!(padded_bbox.min_vertex(), Point2::new(0.5, 1.5)); + assert_eq!(padded_bbox.size(), Vector2::new(4.0, 5.0)); +} + +#[test] +fn test_bounding_box_scaling_2d() { + use nalgebra::{Point2, Vector2}; + + let point = Point2::new(1.0, 1.0); + let size = Vector2::new(3.0, 4.0); + let bbox = AxisAlignedBoundingBox::new(point, size); + + let padded_bbox = bbox.scale(Vector2::new(2.0, 2.0)); + assert_eq!(padded_bbox.min_vertex(), Point2::new(-2.0, -3.0)); + assert_eq!(padded_bbox.size(), Vector2::new(6.0, 8.0)); +} + +#[test] +fn test_bounding_box_contains_2d() { + use nalgebra::Point2; + + let point1 = Point2::new(1.0, 2.0); + let point2 = Point2::new(4.0, 6.0); + let bbox = AxisAlignedBoundingBox::new_2d(point1, point2); + + assert!(bbox.contains_point(&Point2::new(2.0, 3.0))); + assert!(!bbox.contains_point(&Point2::new(5.0, 7.0))); +} + +#[test] +fn test_bounding_box_union_2d() { + use nalgebra::{Point2, Vector2}; + + let point1 = Point2::new(1.0, 2.0); + let point2 = Point2::new(4.0, 6.0); + let bbox1 = AxisAlignedBoundingBox::new_2d(point1, point2); + + let point3 = Point2::new(3.0, 5.0); + let point4 = Point2::new(7.0, 8.0); + let bbox2 = AxisAlignedBoundingBox::new_2d(point3, point4); + + let union_bbox = bbox1.union(&bbox2); + assert_eq!(union_bbox.min_vertex(), Point2::new(1.0, 2.0)); + assert_eq!(union_bbox.size(), Vector2::new(6.0, 6.0)); +} + +#[test] +fn test_bounding_box_intersection_2d() { + use nalgebra::{Point2, Vector2}; + + let point1 = Point2::new(1.0, 2.0); + let point2 = Point2::new(4.0, 6.0); + let bbox1 = AxisAlignedBoundingBox::new_2d(point1, point2); + + let point3 = Point2::new(3.0, 5.0); + let point4 = Point2::new(5.0, 7.0); + let bbox2 = AxisAlignedBoundingBox::new_2d(point3, point4); + + let intersection_bbox = bbox1.intersection(&bbox2).unwrap(); + assert_eq!(intersection_bbox.min_vertex(), Point2::new(3.0, 5.0)); + assert_eq!(intersection_bbox.size(), Vector2::new(1.0, 1.0)); } diff --git a/bounding-box/src/nms.rs b/bounding-box/src/nms.rs new file mode 100644 index 0000000..fe0dc5b --- /dev/null +++ b/bounding-box/src/nms.rs @@ -0,0 +1,68 @@ +use std::collections::HashSet; + +use crate::*; +/// Apply Non-Maximum Suppression to a set of bounding boxes. +/// +/// # Arguments +/// +/// * `boxes` - A slice of bounding boxes to apply NMS on. +/// * `threshold` - The IoU threshold for suppression. +/// +/// # Returns +/// +/// A vector of indices of the bounding boxes that are kept after applying NMS. +pub fn nms( + boxes: &[Aabb2], + scores: &[T], + score_threshold: T, + nms_threshold: T, +) -> Vec> +where + T: Num + num::Float + core::iter::Product + core::ops::AddAssign + core::ops::SubAssign, +{ + use itertools::Itertools; + let bboxes: Vec<_> = boxes + .iter() + .zip(scores.iter()) + .filter_map(|(bbox, score)| (score >= &score_threshold).then_some((bbox, score))) + .sorted_by(|(_, score_a), (_, score_b)| { + score_b + .partial_cmp(score_a) + .unwrap_or(std::cmp::Ordering::Equal) + }) + .map(|(bbox, _)| bbox) + .collect(); + + let outputs = bboxes + .iter() + .enumerate() + .scan( + HashSet::with_capacity(bboxes.len()), + |state, (index, bbox)| { + if state.is_empty() { + state.insert(index); + return Some(Some(bbox)); + } else { + if state.contains(&index) { + return Some(None); + } + let to_remove = bboxes + .iter() + .enumerate() + .skip(index + 1) + .filter_map(|(index, bbox_b)| { + (!state.contains(&index)).then_some(index)?; + let iou = bbox.iou(bbox_b)?; + (iou >= nms_threshold).then_some(index) + }) + .collect_vec(); + state.extend(to_remove); + Some(Some(bbox)) + } + }, + ) + .flatten() + .map(|bbox| **bbox) + .collect_vec(); + outputs +} diff --git a/bounding-box/target/rust-analyzer/metadata/sysroot/Cargo.lock b/bounding-box/target/rust-analyzer/metadata/sysroot/Cargo.lock index 6b1a0a0..5100b4d 100644 --- a/bounding-box/target/rust-analyzer/metadata/sysroot/Cargo.lock +++ b/bounding-box/target/rust-analyzer/metadata/sysroot/Cargo.lock @@ -32,12 +32,6 @@ dependencies = [ "core", ] -[[package]] -name = "allocator-api2" -version = "0.2.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" - [[package]] name = "alloctests" version = "0.0.0" @@ -67,9 +61,9 @@ dependencies = [ [[package]] name = "compiler_builtins" -version = "0.1.152" +version = "0.1.158" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2153cf213eb259361567720ce55f6446f17acd0ccca87fb6dc05360578228a58" +checksum = "164cdc689e4c6d69417f77a5f48be240c291e84fbef0b1281755dc754b19c809" dependencies = [ "cc", "rustc-std-workspace-core", @@ -89,9 +83,9 @@ dependencies = [ [[package]] name = "dlmalloc" -version = "0.2.7" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9b5e0d321d61de16390ed273b647ce51605b575916d3c25e6ddf27a1e140035" +checksum = "8cff88b751e7a276c4ab0e222c3f355190adc6dde9ce39c851db39da34990df7" dependencies = [ "cfg-if", "compiler_builtins", @@ -134,11 +128,10 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.15.2" +version = "0.15.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +checksum = "84b26c544d002229e640969970a2e74021aadf6e2f96372b9c58eff97de08eb3" dependencies = [ - "allocator-api2", "compiler_builtins", "rustc-std-workspace-alloc", "rustc-std-workspace-core", @@ -157,9 +150,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.171" +version = "0.2.172" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" +checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa" dependencies = [ "rustc-std-workspace-core", ] @@ -176,9 +169,9 @@ dependencies = [ [[package]] name = "miniz_oxide" -version = "0.8.3" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8402cab7aefae129c6977bb0ff1b8fd9a04eb5b51efc50a70bea51cda0c7924" +checksum = "3be647b768db090acb35d5ec5db2b0e1f1de11133ca123b9eacf5137868f892a" dependencies = [ "adler2", "compiler_builtins", @@ -221,20 +214,12 @@ dependencies = [ "unwind", ] -[[package]] -name = "proc-macro2" -version = "1.0.93" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99" -dependencies = [ - "unicode-ident", -] - [[package]] name = "proc_macro" version = "0.0.0" dependencies = [ "core", + "rustc-literal-escaper", "std", ] @@ -245,20 +230,11 @@ dependencies = [ "cc", ] -[[package]] -name = "quote" -version = "1.0.38" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc" -dependencies = [ - "proc-macro2", -] - [[package]] name = "r-efi" -version = "4.5.0" +version = "5.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9e935efc5854715dfc0a4c9ef18dc69dee0ec3bf9cc3ab740db831c0fdd86a3" +checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5" dependencies = [ "compiler_builtins", "rustc-std-workspace-core", @@ -266,9 +242,9 @@ dependencies = [ [[package]] name = "r-efi-alloc" -version = "1.0.0" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31d6f09fe2b6ad044bc3d2c34ce4979796581afd2f1ebc185837e02421e02fd7" +checksum = "e43c53ff1a01d423d1cb762fd991de07d32965ff0ca2e4f80444ac7804198203" dependencies = [ "compiler_builtins", "r-efi", @@ -277,22 +253,18 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.0" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" +checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" dependencies = [ "rand_core", - "zerocopy", ] [[package]] name = "rand_core" -version = "0.9.0" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff" -dependencies = [ - "zerocopy", -] +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" [[package]] name = "rand_xorshift" @@ -313,6 +285,15 @@ dependencies = [ "rustc-std-workspace-core", ] +[[package]] +name = "rustc-literal-escaper" +version = "0.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0041b6238913c41fe704213a4a9329e2f685a156d1781998128b4149c230ad04" +dependencies = [ + "rustc-std-workspace-std", +] + [[package]] name = "rustc-std-workspace-alloc" version = "1.99.0" @@ -380,17 +361,6 @@ dependencies = [ "rustc-std-workspace-core", ] -[[package]] -name = "syn" -version = "2.0.98" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - [[package]] name = "sysroot" version = "0.0.0" @@ -411,12 +381,6 @@ dependencies = [ "std", ] -[[package]] -name = "unicode-ident" -version = "1.0.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a210d160f08b701c8721ba1c726c11662f877ea6b7094007e1ca9a1041945034" - [[package]] name = "unicode-width" version = "0.1.14" @@ -441,9 +405,9 @@ dependencies = [ [[package]] name = "unwinding" -version = "0.2.5" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51f06a05848f650946acef3bf525fe96612226b61f74ae23ffa4e98bfbb8ab3c" +checksum = "8393f2782b6060a807337ff353780c1ca15206f9ba2424df18cb6e733bd7b345" dependencies = [ "compiler_builtins", "gimli", @@ -537,23 +501,3 @@ name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" - -[[package]] -name = "zerocopy" -version = "0.8.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa91407dacce3a68c56de03abe2760159582b846c6a4acd2f456618087f12713" -dependencies = [ - "zerocopy-derive", -] - -[[package]] -name = "zerocopy-derive" -version = "0.8.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06718a168365cad3d5ff0bb133aad346959a2074bd4a85c121255a11304a8626" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] diff --git a/flake.nix b/flake.nix index d69dd7f..ddd2e70 100644 --- a/flake.nix +++ b/flake.nix @@ -83,7 +83,7 @@ pname = name; stdenv = pkgs.clangStdenv; doCheck = false; - # LIBCLANG_PATH = "${pkgs.llvmPackages.libclang.lib}/lib"; + LIBCLANG_PATH = "${pkgs.llvmPackages.libclang.lib}/lib"; # nativeBuildInputs = with pkgs; [ # cmake # llvmPackages.libclang.lib @@ -160,6 +160,7 @@ stableToolchainWithRustAnalyzer cargo-nextest cargo-deny + cmake mnn ] ++ (lib.optionals pkgs.stdenv.isDarwin [ diff --git a/src/facedet.rs b/src/facedet.rs index 916f5a6..214ad33 100644 --- a/src/facedet.rs +++ b/src/facedet.rs @@ -1,13 +1,14 @@ use crate::errors::*; +use bounding_box::Aabb2; use error_stack::ResultExt; use mnn_bridge::ndarray::*; -use nalgebra::Point2; +use nalgebra::{Point2, Vector2}; use ndarray_resize::NdFir; use std::path::Path; pub struct FaceDetectionConfig { - min_sizes: Vec>, - steps: Vec, + min_sizes: Vec>, + steps: Vec, variance: Vec, } pub struct FaceDetection { @@ -20,6 +21,37 @@ pub struct FaceDetectionModelOutput { pub landmark: ndarray::Array3, } +impl FaceDetectionModelOutput { + pub fn postprocess(self, config: FaceDetectionConfig) -> Result>> { + // for k, step in enumerate(cfg['steps']): + // feature_size = 640 // step + // for i in range(feature_size): + // for j in range(feature_size): + // for min_size in cfg['min_sizes'][k]: + // cx = (j + 0.5) * step / 640 + // cy = (i + 0.5) * step / 640 + // s_kx = s_ky = min_size / 640 + // anchors.append([cx, cy, s_kx, s_ky]) + let mut anchors = Vec::new(); + config.steps.iter().enumerate().for_each(|(k, step)| { + let feature_size = 640 / step; + for i in 0..feature_size { + for j in 0..feature_size { + for min_size in &config.min_sizes[k] { + let cx = (j as f32 + 0.5) * *step as f32 / 640.0; + let cy = (i as f32 + 0.5) * *step as f32 / 640.0; + let s_kx = *min_size as f32 / 640.0; + let s_ky = *min_size as f32 / 640.0; + anchors.push([cx, cy, s_kx, s_ky]); + } + } + } + }); + + Ok(Vec::new()) + } +} + impl FaceDetectionModelOutput { pub fn print(&self, limit: usize) { tracing::info!("Detected {} faces", self.bbox.shape()[1]);