servius/face-detector
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Added a manual implementation of nms

Browse Source
This commit is contained in:
uttarayan21
2025-08-07 15:45:54 +05:30
parent 2c43f657aa
commit e91ae5b865
5 changed files with 46 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@@ -255,6 +255,7 @@ dependencies = [
"nalgebra",
"ndarray 0.16.1",
"num",
"ordered-float",
"simba",
"thiserror 2.0.12",
]

1
bounding-box/Cargo.toml
View File

@@ -9,6 +9,7 @@ itertools = "0.14.0"
nalgebra = "0.33.2"
ndarray = { version = "0.16.1", optional = true }
num = "0.4.3"
ordered-float = "5.0.0"
simba = "0.9.0"
thiserror = "2.0.12"

27
bounding-box/src/lib.rs
View File

@@ -51,10 +51,12 @@ pub type Aabb3<T> = AxisAlignedBoundingBox<T, 3>;
impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
// Panics if max < min
pub fn new(min_point: Point<T, D>, max_point: Point<T, D>) -> Self {
if max_point < min_point {
if max_point >= min_point {
Self::from_min_max_vertices(min_point, max_point)
} else {
dbg!(max_point, min_point);
panic!("max_point must be greater than or equal to min_point");
}
Self::from_min_max_vertices(min_point, max_point)
}
pub fn try_new(min_point: Point<T, D>, max_point: Point<T, D>) -> Option<Self> {
if max_point < min_point {
@@ -66,9 +68,9 @@ impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
Self { point, size }
}
pub fn from_min_max_vertices(point1: Point<T, D>, point2: Point<T, D>) -> Self {
let size = point2 - point1;
Self::new_point_size(point1, SVector::from(size))
pub fn from_min_max_vertices(min: Point<T, D>, max: Point<T, D>) -> Self {
let size = max - min;
Self::new_point_size(min, SVector::from(size))
}
/// Only considers the points closest and furthest from origin
@@ -301,11 +303,11 @@ impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
let inter_min = lhs_min.sup(&rhs_min);
let inter_max = lhs_max.inf(&rhs_max);
if inter_max < inter_min {
return T::zero();
} else {
if inter_max >= inter_min {
let intersection = Aabb::new(inter_min, inter_max).measure();
intersection / (self.measure() + other.measure() - intersection)
} else {
return T::zero();
}
}
}
@@ -605,11 +607,8 @@ mod boudning_box_tests {
#[test]
fn test_specific_values() {
let res = Vector2::new(1920, 1080).cast();
let box1 = Aabb2::from_xywh(0.69482, 0.6716774, 0.07493961, 0.14968264).denormalize(res);
let box2 =
Aabb2::from_xywh(0.41546485, 0.70290875, 0.06197411, 0.08818436).denormalize(res);
dbg!(box1, box2);
assert!(box1.iou(&box2) > 0.0);
let box1 = Aabb2::from_xywh(0.69482, 0.6716774, 0.07493961, 0.14968264);
let box2 = Aabb2::from_xywh(0.41546485, 0.70290875, 0.06197411, 0.08818436);
assert!(box1.iou(&box2) >= 0.0);
}
}

93
bounding-box/src/nms.rs
View File

@@ -1,4 +1,6 @@
use std::collections::HashSet;
use std::collections::{HashSet, VecDeque};
use itertools::Itertools;
use crate::*;
/// Apply Non-Maximum Suppression to a set of bounding boxes.
@@ -21,7 +23,8 @@ pub fn nms<T>(
) -> HashSet<usize>
where
T: Num
+ num::Float
+ ordered_float::FloatCore
+ core::ops::Neg<Output = T>
+ core::iter::Product<T>
+ core::ops::AddAssign
+ core::ops::SubAssign
@@ -29,71 +32,31 @@ where
+ nalgebra::SimdValue
+ nalgebra::SimdPartialOrd,
{
let mut indices: Vec<usize> = (0..boxes.len())
.filter(|&i| scores[i] >= score_threshold)
let mut combined: VecDeque<(usize, Aabb2<T>, T, bool)> = boxes
.iter()
.enumerate()
.zip(scores)
.filter_map(|((idx, bbox), score)| {
(*score > score_threshold).then_some((idx, *bbox, *score, true))
})
.sorted_by_cached_key(|(_, _, score, _)| -ordered_float::OrderedFloat(*score))
.collect();
indices.sort_by(|&i, &j| scores[j].partial_cmp(&scores[i]).unwrap());
let mut selected_indices = HashSet::new();
while let Some(&current) = indices.first() {
selected_indices.insert(current);
indices.remove(0);
indices.retain(|&i| {
let iou = calculate_iou(&boxes[current], &boxes[i]);
let iou_ = boxes[current].iou(&boxes[i]);
if iou != iou_ {
dbg!(boxes[current], boxes[i]);
panic!()
};
iou < nms_threshold
});
for i in 0..combined.len() {
let first = combined[i];
if first.3 == false {
continue;
}
selected_indices
}
/// Calculate the Intersection over Union (IoU) of two bounding boxes.
///
/// # Arguments
///
/// * `box1` - The first bounding box.
/// * `box2` - The second bounding box.
///
/// # Returns
///
/// The IoU as a value between 0 and 1.
fn calculate_iou<T>(box1: &Aabb2<T>, box2: &Aabb2<T>) -> T
where
T: Num
+ num::Float
+ core::iter::Product<T>
+ core::ops::AddAssign
+ core::ops::SubAssign
+ core::ops::MulAssign
+ nalgebra::SimdValue
+ nalgebra::SimdPartialOrd,
{
let x_left = box1.min_vertex().x.max(box2.min_vertex().x);
let y_top = box1.min_vertex().y.max(box2.min_vertex().y);
let x_right = box1.max_vertex().x.min(box2.max_vertex().x);
let y_bottom = box1.max_vertex().y.min(box2.max_vertex().y);
let zero = T::zero();
let inter_width = (x_right - x_left).max(zero);
let inter_height = (y_bottom - y_top).max(zero);
let intersection = inter_width * inter_height;
let area1 = box1.area();
let area2 = box2.area();
let union = area1 + area2 - intersection;
if union > zero {
intersection / union
} else {
zero
let bbox = first.1;
for item in combined.iter_mut().skip(i + 1) {
if bbox.iou(&item.1) > nms_threshold {
item.3 = false
}
}
}
combined
.into_iter()
.filter_map(|(idx, _, _, keep)| keep.then_some(idx))
.collect()
}

4
src/main.rs
View File

@@ -27,7 +27,9 @@ pub fn main() -> Result<()> {
let output = model
.detect_faces(
array.clone(),
FaceDetectionConfig::default().with_threshold(detect.threshold),
FaceDetectionConfig::default()
.with_threshold(detect.threshold)
.with_nms_threshold(detect.nms_threshold),
)
.change_context(errors::Error)
.attach_printable("Failed to detect faces")?;