feat: Remove bbox crate and use 1024 for image size

bounding-box/src/draw.rs

@@ -53,8 +53,9 @@ impl Drawable<Array3<u8>> for Aabb2<usize> {
         let bottom = Aabb2::from_x1y1x2y2(x1y2.x, x1y2.y, x2y2.x, x2y2.y + thickness);
         let left = Aabb2::from_x1y1x2y2(x1y1.x, x1y1.y, x1y2.x + thickness, x1y2.y);
         let right = Aabb2::from_x1y1x2y2(x2y1.x, x2y1.y, x2y2.x + thickness, x2y2.y + thickness);
-        let lines = [top, bottom, left, right];
-        lines.into_iter().for_each(|line| {
+        let canvas_bbox = Aabb2::from_x1y1x2y2(0, 0, canvas.dim().1 - 1, canvas.dim().0 - 1);
+        let lines = [top, bottom, left, right].map(|bbox| bbox.clamp(&canvas_bbox));
+        lines.into_iter().flatten().for_each(|line| {
             canvas
                 .roi_mut(line)
                 .map(|mut line| {

bounding-box/src/nms.rs

@@ -6,7 +6,9 @@ use crate::*;
 /// # Arguments
 ///
 /// * `boxes` - A slice of bounding boxes to apply NMS on.
-/// * `threshold` - The IoU threshold for suppression.
+/// * `scores` - A slice of confidence scores corresponding to the bounding boxes.
+/// * `score_threshold` - The minimum score threshold for consideration.
+/// * `nms_threshold` - The IoU threshold for suppression.
 ///
 /// # Returns
 ///
@@ -16,7 +18,7 @@ pub fn nms<T>(
     scores: &[T],
     score_threshold: T,
     nms_threshold: T,
-) -> Vec<Aabb2<T>>
+) -> HashSet<usize>
 where
     T: Num
         + num::Float
@@ -28,48 +30,55 @@ where
         + nalgebra::SimdPartialOrd,
 {
     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
+    // Create vector of (index, box, score) tuples for boxes with scores above threshold
+    let mut indexed_boxes: Vec<(usize, &Aabb2<T>, &T)> = boxes
         .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))
+        .zip(scores.iter())
+        .filter_map(|((idx, bbox), score)| {
+            if *score >= score_threshold {
+                Some((idx, bbox, score))
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    // Sort by score in descending order
+    indexed_boxes.sort_by(|(_, _, score_a), (_, _, score_b)| {
+        score_b
+            .partial_cmp(score_a)
+            .unwrap_or(std::cmp::Ordering::Equal)
+    });
+
+    let mut keep_indices = HashSet::new();
+    let mut suppressed = HashSet::new();
+
+    for (i, (idx_i, bbox_i, _)) in indexed_boxes.iter().enumerate() {
+        // Skip if this box is already suppressed
+        if suppressed.contains(idx_i) {
+            continue;
+        }
+
+        // Keep this box
+        keep_indices.insert(*idx_i);
+
+        // Compare with remaining boxes
+        for (idx_j, bbox_j, _) in indexed_boxes.iter().skip(i + 1) {
+            // Skip if this box is already suppressed
+            if suppressed.contains(idx_j) {
+                continue;
+            }
+
+            // Calculate IoU and suppress if above threshold
+            if let Some(iou) = bbox_i.iou(bbox_j) {
+                if iou >= nms_threshold {
+                    suppressed.insert(*idx_j);
                 }
-            },
-        )
-        .flatten()
-        .map(|bbox| **bbox)
-        .collect_vec();
-    outputs
+            }
+        }
+    }
+
+    keep_indices
 }