@camera.ui/rust-postprocessor 0.0.1

package/src/nms.rs

@@ -0,0 +1,296 @@
+//! Non-Maximum Suppression with SIMD acceleration.
+//!
+//! Greedy NMS over a list of detections, suppressing within the same class
+//! label only. Uses a Structure-of-Arrays layout so the inner IoU loop can
+//! process eight candidate boxes at once via `wide::f32x8`.
+//!
+//! Algorithm steps:
+//!   1. Optional top-k pre-filter — bounds the work for pathological inputs
+//!      with thousands of overlapping detections.
+//!   2. Sort surviving candidates by descending confidence.
+//!   3. Greedy keep-then-suppress: for each kept box compare in `f32x8`
+//!      chunks against all later boxes of the same class, suppressing those
+//!      whose IoU exceeds the threshold.
+//!
+//! Output preserves the original `Detection` order by confidence (descending).
+
+use wide::{f32x8, CmpGt};
+// CmpGt is the trait that provides `simd_gt` on f32x8.
+
+use crate::types::Detection;
+
+/// Run greedy NMS on `detections`.
+///
+/// * `iou_threshold` — boxes with IoU strictly greater than this value
+///   compared to a higher-scoring box of the same class are removed.
+/// * `max_detections` — hard cap on the output length. `None` for no cap.
+///
+/// Detections without overlapping rivals pass through unchanged. Per-class
+/// semantics ensure a `person` and a `car` with overlapping bboxes are both
+/// retained.
+pub fn nms(
+  detections: Vec<Detection>,
+  iou_threshold: f32,
+  max_detections: Option<usize>,
+) -> Vec<Detection> {
+  if detections.is_empty() {
+    return Vec::new();
+  }
+
+  // Pre-filter: keep at most max_detections * 10 by confidence to bound the
+  // O(n²) inner loop on adversarial inputs. The factor 10 leaves headroom for
+  // suppression to still hit the requested final count.
+  let mut candidates = detections;
+  let max_det = max_detections.unwrap_or(candidates.len());
+  let prefilter_cap = max_det.saturating_mul(10).max(64);
+  if candidates.len() > prefilter_cap {
+    candidates.select_nth_unstable_by(prefilter_cap, |a, b| {
+      b.confidence
+        .partial_cmp(&a.confidence)
+        .unwrap_or(std::cmp::Ordering::Equal)
+    });
+    candidates.truncate(prefilter_cap);
+  }
+
+  // Sort by confidence descending so the highest-score box is always picked
+  // first in the greedy pass.
+  candidates.sort_unstable_by(|a, b| {
+    b.confidence
+      .partial_cmp(&a.confidence)
+      .unwrap_or(std::cmp::Ordering::Equal)
+  });
+
+  let n = candidates.len();
+
+  // Structure-of-Arrays for SIMD-friendly access. We store xyxy (corners)
+  // because the IoU intersection math uses min/max over corners directly.
+  let mut x1 = vec![0.0f32; n];
+  let mut y1 = vec![0.0f32; n];
+  let mut x2 = vec![0.0f32; n];
+  let mut y2 = vec![0.0f32; n];
+  let mut areas = vec![0.0f32; n];
+
+  for (i, det) in candidates.iter().enumerate() {
+    x1[i] = det.x;
+    y1[i] = det.y;
+    x2[i] = det.x + det.width;
+    y2[i] = det.y + det.height;
+    areas[i] = det.width * det.height;
+  }
+
+  // Per-class label dispatch: convert label strings to small integer ids so
+  // the inner loop can compare classes via cheap integer equality.
+  let mut label_ids = vec![0u32; n];
+  let mut label_table: Vec<&str> = Vec::new();
+  for (i, det) in candidates.iter().enumerate() {
+    let id = match label_table.iter().position(|l| *l == det.label.as_str()) {
+      Some(idx) => idx as u32,
+      None => {
+        label_table.push(det.label.as_str());
+        (label_table.len() - 1) as u32
+      }
+    };
+    label_ids[i] = id;
+  }
+
+  let mut suppressed = vec![false; n];
+  let mut keep_indices: Vec<usize> = Vec::with_capacity(max_det.min(n));
+
+  let iou_v = f32x8::splat(iou_threshold);
+  let zero_v = f32x8::splat(0.0);
+
+  for i in 0..n {
+    if suppressed[i] {
+      continue;
+    }
+    keep_indices.push(i);
+    if keep_indices.len() >= max_det {
+      break;
+    }
+
+    let ax1 = f32x8::splat(x1[i]);
+    let ay1 = f32x8::splat(y1[i]);
+    let ax2 = f32x8::splat(x2[i]);
+    let ay2 = f32x8::splat(y2[i]);
+    let aa = f32x8::splat(areas[i]);
+    let a_label = label_ids[i];
+
+    let mut j = i + 1;
+    // SIMD chunks of 8 — load 8 boxes at a time, compute IoU vector, build
+    // a suppression mask, then apply it back per-element with class check.
+    while j + 8 <= n {
+      // Skip the chunk entirely if no candidate inside it can be suppressed
+      // (all already suppressed OR all of a different class).
+      let mut chunk_active = false;
+      for k in 0..8 {
+        if !suppressed[j + k] && label_ids[j + k] == a_label {
+          chunk_active = true;
+          break;
+        }
+      }
+
+      if chunk_active {
+        // Safe unaligned reads — Vec<f32> is contiguous and we never index
+        // past the end (the `while` condition guarantees j + 8 <= n).
+        let bx1: f32x8 = unsafe { (x1.as_ptr().add(j) as *const f32x8).read_unaligned() };
+        let by1: f32x8 = unsafe { (y1.as_ptr().add(j) as *const f32x8).read_unaligned() };
+        let bx2: f32x8 = unsafe { (x2.as_ptr().add(j) as *const f32x8).read_unaligned() };
+        let by2: f32x8 = unsafe { (y2.as_ptr().add(j) as *const f32x8).read_unaligned() };
+        let ba: f32x8 = unsafe { (areas.as_ptr().add(j) as *const f32x8).read_unaligned() };
+
+        let ix1 = ax1.fast_max(bx1);
+        let iy1 = ay1.fast_max(by1);
+        let ix2 = ax2.fast_min(bx2);
+        let iy2 = ay2.fast_min(by2);
+
+        let iw = (ix2 - ix1).fast_max(zero_v);
+        let ih = (iy2 - iy1).fast_max(zero_v);
+        let inter = iw * ih;
+        let union = aa + ba - inter;
+        // Avoid divide-by-zero — when union is 0 the IoU is irrelevant
+        // (both boxes have zero area, the candidate would be filtered
+        // elsewhere). Use a sentinel large enough to fail the threshold
+        // check below.
+        let safe_union = union.fast_max(f32x8::splat(f32::MIN_POSITIVE));
+        let iou = inter / safe_union;
+
+        let mask = iou.simd_gt(iou_v);
+        // `to_bitmask` packs the SIMD comparison result into a bitmap where
+        // bit k is set when lane k of `iou` exceeds the threshold.
+        let bits = mask.to_bitmask();
+        if bits != 0 {
+          for k in 0..8 {
+            if (bits & (1 << k)) != 0 && !suppressed[j + k] && label_ids[j + k] == a_label {
+              suppressed[j + k] = true;
+            }
+          }
+        }
+      }
+      j += 8;
+    }
+
+    // Scalar tail — handle the remaining < 8 elements without SIMD.
+    while j < n {
+      if !suppressed[j] && label_ids[j] == a_label {
+        let ix1 = x1[i].max(x1[j]);
+        let iy1 = y1[i].max(y1[j]);
+        let ix2 = x2[i].min(x2[j]);
+        let iy2 = y2[i].min(y2[j]);
+        let iw = (ix2 - ix1).max(0.0);
+        let ih = (iy2 - iy1).max(0.0);
+        let inter = iw * ih;
+        let union = areas[i] + areas[j] - inter;
+        if union > 0.0 && inter / union > iou_threshold {
+          suppressed[j] = true;
+        }
+      }
+      j += 1;
+    }
+  }
+
+  // Move kept detections out of the candidates Vec — no clones, ownership
+  // transfer via take() over a Vec<Option<Detection>> would also work but
+  // an indexed swap_remove pattern is simpler since keep_indices is sorted.
+  let mut kept: Vec<Detection> = Vec::with_capacity(keep_indices.len());
+  // Build a flag vector so we can drain in original order with one pass.
+  let mut keep_flag = vec![false; n];
+  for &idx in &keep_indices {
+    keep_flag[idx] = true;
+  }
+  for (i, det) in candidates.into_iter().enumerate() {
+    if keep_flag[i] {
+      kept.push(det);
+    }
+  }
+  kept
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  fn det(x: f32, y: f32, w: f32, h: f32, conf: f32, label: &str) -> Detection {
+    Detection {
+      x,
+      y,
+      width: w,
+      height: h,
+      confidence: conf,
+      label: label.to_string(),
+    }
+  }
+
+  #[test]
+  fn empty_input() {
+    assert!(nms(Vec::new(), 0.5, None).is_empty());
+  }
+
+  #[test]
+  fn single_box_kept() {
+    let input = vec![det(0.1, 0.1, 0.2, 0.2, 0.9, "person")];
+    let out = nms(input, 0.5, None);
+    assert_eq!(out.len(), 1);
+  }
+
+  #[test]
+  fn duplicate_boxes_suppressed() {
+    let input = vec![
+      det(0.1, 0.1, 0.2, 0.2, 0.9, "person"),
+      det(0.1, 0.1, 0.2, 0.2, 0.8, "person"),
+      det(0.1, 0.1, 0.2, 0.2, 0.7, "person"),
+    ];
+    let out = nms(input, 0.5, None);
+    assert_eq!(out.len(), 1);
+    assert!((out[0].confidence - 0.9).abs() < 1e-6);
+  }
+
+  #[test]
+  fn different_classes_kept() {
+    let input = vec![
+      det(0.1, 0.1, 0.2, 0.2, 0.9, "person"),
+      det(0.1, 0.1, 0.2, 0.2, 0.8, "car"),
+    ];
+    let out = nms(input, 0.5, None);
+    assert_eq!(out.len(), 2);
+  }
+
+  #[test]
+  fn many_overlapping_simd_path() {
+    // Generate 20 nearly-identical person boxes — exercises the SIMD
+    // chunked loop (n - 1 = 19, so multiple f32x8 iterations).
+    let mut input = Vec::new();
+    for i in 0..20 {
+      input.push(det(0.1, 0.1, 0.2, 0.2, 0.9 - i as f32 * 0.001, "person"));
+    }
+    let out = nms(input, 0.5, None);
+    assert_eq!(out.len(), 1);
+  }
+
+  #[test]
+  fn max_detections_cap() {
+    let input = vec![
+      det(0.0, 0.0, 0.1, 0.1, 0.9, "a"),
+      det(0.2, 0.2, 0.1, 0.1, 0.85, "b"),
+      det(0.4, 0.4, 0.1, 0.1, 0.8, "c"),
+      det(0.6, 0.6, 0.1, 0.1, 0.75, "d"),
+    ];
+    let out = nms(input, 0.5, Some(2));
+    assert_eq!(out.len(), 2);
+    // Top two by confidence
+    assert!((out[0].confidence - 0.9).abs() < 1e-6);
+    assert!((out[1].confidence - 0.85).abs() < 1e-6);
+  }
+
+  #[test]
+  fn output_sorted_by_confidence() {
+    let input = vec![
+      det(0.0, 0.0, 0.1, 0.1, 0.5, "a"),
+      det(0.2, 0.2, 0.1, 0.1, 0.9, "b"),
+      det(0.4, 0.4, 0.1, 0.1, 0.7, "c"),
+    ];
+    let out = nms(input, 0.5, None);
+    assert_eq!(out.len(), 3);
+    assert!(out[0].confidence > out[1].confidence);
+    assert!(out[1].confidence > out[2].confidence);
+  }
+}