sightloom_core/
nms.rs

1//! Deterministic non-maximum suppression over caller-owned storage.
2
3use core::cmp::Ordering;
4
5use crate::{CoreError, Detection, ios, iou};
6
7/// The rectangle-overlap metric used to suppress detections.
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9pub enum OverlapMetric {
10    /// Intersection over union.
11    IoU,
12    /// Intersection over the smaller rectangle's area.
13    IoS,
14}
15
16/// The class-comparison policy used during suppression.
17#[derive(Clone, Copy, Debug, Eq, PartialEq)]
18pub enum NmsMode {
19    /// Compare detections only when their optional class identifiers match.
20    ClassAware,
21    /// Compare every detection regardless of class identifier.
22    ClassAgnostic,
23}
24
25/// Immutable settings for a non-maximum suppression pass.
26#[derive(Clone, Copy, Debug, PartialEq)]
27pub struct NmsConfig {
28    /// Overlap threshold in the inclusive range `0.0..=1.0`; suppression
29    /// occurs only when overlap is strictly greater than this value.
30    pub threshold: f32,
31    /// Whether class identifiers restrict suppression.
32    pub mode: NmsMode,
33    /// The overlap metric used for suppression.
34    pub metric: OverlapMetric,
35}
36
37/// Suppresses overlapping detections without allocating.
38///
39/// Detections are prioritized by descending score, then ascending original
40/// index. The kept detections are compacted into the front of `detections` in
41/// their original input order; the remaining tail is unspecified.
42///
43/// # Errors
44///
45/// Returns [`CoreError::InvalidThreshold`] when `config.threshold` is not
46/// finite or is outside `0.0..=1.0`. Returns
47/// [`CoreError::InsufficientScratch`] when either scratch slice is shorter
48/// than `detections`. Neither error mutates a caller-owned slice.
49pub fn nms_in_place(
50    detections: &mut [Detection],
51    order_scratch: &mut [usize],
52    suppressed_scratch: &mut [bool],
53    config: NmsConfig,
54) -> Result<usize, CoreError> {
55    if !config.threshold.is_finite() || !(0.0..=1.0).contains(&config.threshold) {
56        return Err(CoreError::InvalidThreshold);
57    }
58
59    let len = detections.len();
60    if order_scratch.len() < len || suppressed_scratch.len() < len {
61        return Err(CoreError::InsufficientScratch);
62    }
63
64    let order = &mut order_scratch[..len];
65    let suppressed = &mut suppressed_scratch[..len];
66    for (index, slot) in order.iter_mut().enumerate() {
67        *slot = index;
68    }
69    for value in suppressed.iter_mut() {
70        *value = false;
71    }
72
73    order.sort_unstable_by(|left, right| {
74        detections[*right]
75            .score()
76            .partial_cmp(&detections[*left].score())
77            .unwrap_or(Ordering::Equal)
78            .then_with(|| left.cmp(right))
79    });
80
81    for candidate_position in 0..len {
82        let candidate_index = order[candidate_position];
83        if suppressed[candidate_index] {
84            continue;
85        }
86
87        for contender_index in order.iter().skip(candidate_position + 1).copied() {
88            if suppressed[contender_index]
89                || (config.mode == NmsMode::ClassAware
90                    && detections[candidate_index].class_id()
91                        != detections[contender_index].class_id())
92            {
93                continue;
94            }
95
96            let overlap = match config.metric {
97                OverlapMetric::IoU => iou(
98                    detections[candidate_index].bbox(),
99                    detections[contender_index].bbox(),
100                ),
101                OverlapMetric::IoS => ios(
102                    detections[candidate_index].bbox(),
103                    detections[contender_index].bbox(),
104                ),
105            };
106            if overlap > config.threshold {
107                suppressed[contender_index] = true;
108            }
109        }
110    }
111
112    let mut kept = 0;
113    for index in 0..len {
114        if !suppressed[index] {
115            let detection = detections[index];
116            detections[kept] = detection;
117            kept += 1;
118        }
119    }
120
121    Ok(kept)
122}
sightloom_core/nms.rs

sightloom_core/
nms.rs
