spatialrust_vision/
detection.rs

1//! Detection post-processing primitives.
2
3use std::cmp::Ordering;
4use std::collections::HashMap;
5
6use crate::{VisionError, VisionResult};
7
8/// Axis-aligned bounding box using half-open `(min, max)` coordinates.
9#[derive(Clone, Copy, Debug, Default, PartialEq)]
10pub struct BoundingBox2 {
11    /// Minimum x coordinate.
12    pub x_min: f32,
13    /// Minimum y coordinate.
14    pub y_min: f32,
15    /// Maximum x coordinate.
16    pub x_max: f32,
17    /// Maximum y coordinate.
18    pub y_max: f32,
19}
20
21impl BoundingBox2 {
22    /// Creates a validated axis-aligned box.
23    pub fn try_new(x_min: f32, y_min: f32, x_max: f32, y_max: f32) -> VisionResult<Self> {
24        let bbox = Self { x_min, y_min, x_max, y_max };
25        if !bbox.is_valid() {
26            return Err(VisionError::InvalidParameter(
27                "box coordinates must be finite and max >= min".to_owned(),
28            ));
29        }
30        Ok(bbox)
31    }
32
33    /// Returns whether coordinates are finite and ordered.
34    #[must_use]
35    pub fn is_valid(self) -> bool {
36        self.x_min.is_finite()
37            && self.y_min.is_finite()
38            && self.x_max.is_finite()
39            && self.y_max.is_finite()
40            && self.x_max >= self.x_min
41            && self.y_max >= self.y_min
42    }
43
44    /// Box width.
45    #[must_use]
46    pub fn width(self) -> f32 {
47        (self.x_max - self.x_min).max(0.0)
48    }
49
50    /// Box height.
51    #[must_use]
52    pub fn height(self) -> f32 {
53        (self.y_max - self.y_min).max(0.0)
54    }
55
56    /// Box area.
57    #[must_use]
58    pub fn area(self) -> f32 {
59        self.width() * self.height()
60    }
61
62    /// Intersection box, including zero-area edge contact.
63    #[must_use]
64    pub fn intersection(self, other: Self) -> Option<Self> {
65        let result = Self {
66            x_min: self.x_min.max(other.x_min),
67            y_min: self.y_min.max(other.y_min),
68            x_max: self.x_max.min(other.x_max),
69            y_max: self.y_max.min(other.y_max),
70        };
71        result.is_valid().then_some(result)
72    }
73
74    /// Intersection over union, returning zero for two zero-area boxes.
75    #[must_use]
76    pub fn iou(self, other: Self) -> f32 {
77        let intersection = self.intersection(other).map_or(0.0, Self::area);
78        let union = self.area() + other.area() - intersection;
79        if union > 0.0 {
80            intersection / union
81        } else {
82            0.0
83        }
84    }
85
86    /// Generalized intersection over union (GIoU).
87    #[must_use]
88    pub fn generalized_iou(self, other: Self) -> f32 {
89        let iou = self.iou(other);
90        let enclosing = Self {
91            x_min: self.x_min.min(other.x_min),
92            y_min: self.y_min.min(other.y_min),
93            x_max: self.x_max.max(other.x_max),
94            y_max: self.y_max.max(other.y_max),
95        };
96        let enclosing_area = enclosing.area();
97        if enclosing_area == 0.0 {
98            return iou;
99        }
100        let intersection = self.intersection(other).map_or(0.0, Self::area);
101        let union = self.area() + other.area() - intersection;
102        iou - (enclosing_area - union) / enclosing_area
103    }
104
105    /// Clips coordinates into an image rectangle.
106    #[must_use]
107    pub fn clip(self, width: f32, height: f32) -> Self {
108        Self {
109            x_min: self.x_min.clamp(0.0, width),
110            y_min: self.y_min.clamp(0.0, height),
111            x_max: self.x_max.clamp(0.0, width),
112            y_max: self.y_max.clamp(0.0, height),
113        }
114    }
115
116    /// Applies uniform scaling and translation, useful for letterbox mapping.
117    #[must_use]
118    pub fn scale_translate(self, scale: f32, tx: f32, ty: f32) -> Self {
119        Self {
120            x_min: self.x_min.mul_add(scale, tx),
121            y_min: self.y_min.mul_add(scale, ty),
122            x_max: self.x_max.mul_add(scale, tx),
123            y_max: self.y_max.mul_add(scale, ty),
124        }
125    }
126}
127
128/// One scored detection used by class-aware post-processing.
129#[derive(Clone, Copy, Debug, PartialEq)]
130pub struct Detection {
131    /// Detected box.
132    pub bbox: BoundingBox2,
133    /// Confidence score.
134    pub score: f32,
135    /// Model-defined class identifier.
136    pub class_id: i64,
137}
138
139/// Soft-NMS score-decay strategy.
140#[derive(Clone, Copy, Debug, PartialEq)]
141pub enum SoftNmsMethod {
142    /// Hard suppression above the IoU threshold.
143    Hard,
144    /// Linear score decay above the IoU threshold.
145    Linear,
146    /// Gaussian score decay at every overlap.
147    Gaussian {
148        /// Positive Gaussian variance parameter.
149        sigma: f32,
150    },
151}
152
153/// Index and possibly updated score returned by Soft-NMS.
154#[derive(Clone, Copy, Debug, PartialEq)]
155pub struct ScoredIndex {
156    /// Index into the original input arrays.
157    pub index: usize,
158    /// Score after overlap decay.
159    pub score: f32,
160}
161
162/// Greedy non-maximum suppression. Returned indices are score-descending.
163pub fn nms(
164    boxes: &[BoundingBox2],
165    scores: &[f32],
166    score_threshold: f32,
167    iou_threshold: f32,
168) -> VisionResult<Vec<usize>> {
169    validate_nms_inputs(boxes, scores, score_threshold, iou_threshold)?;
170    let mut order: Vec<usize> = scores
171        .iter()
172        .enumerate()
173        .filter_map(|(index, &score)| (score >= score_threshold).then_some(index))
174        .collect();
175    sort_indices_by_score(&mut order, scores);
176    let areas: Vec<f32> = boxes.iter().copied().map(BoundingBox2::area).collect();
177    let mut keep: Vec<usize> = Vec::with_capacity(order.len());
178    'candidate: for index in order {
179        for &selected in &keep {
180            if iou_with_areas(boxes[index], areas[index], boxes[selected], areas[selected])
181                > iou_threshold
182            {
183                continue 'candidate;
184            }
185        }
186        keep.push(index);
187    }
188    Ok(keep)
189}
190
191/// Class-aware greedy NMS over detection records.
192pub fn batched_nms(
193    detections: &[Detection],
194    score_threshold: f32,
195    iou_threshold: f32,
196) -> VisionResult<Vec<usize>> {
197    if !score_threshold.is_finite()
198        || !iou_threshold.is_finite()
199        || !(0.0..=1.0).contains(&iou_threshold)
200    {
201        return Err(VisionError::InvalidParameter(
202            "NMS thresholds must be finite and IoU in [0, 1]".to_owned(),
203        ));
204    }
205    if detections.iter().any(|detection| !detection.bbox.is_valid() || !detection.score.is_finite())
206    {
207        return Err(VisionError::InvalidParameter(
208            "detections must contain valid boxes and finite scores".to_owned(),
209        ));
210    }
211    let scores: Vec<f32> = detections.iter().map(|detection| detection.score).collect();
212    let mut order: Vec<usize> = scores
213        .iter()
214        .enumerate()
215        .filter_map(|(index, &score)| (score >= score_threshold).then_some(index))
216        .collect();
217    sort_indices_by_score(&mut order, &scores);
218    let areas: Vec<f32> = detections.iter().map(|detection| detection.bbox.area()).collect();
219    let mut keep: Vec<usize> = Vec::with_capacity(order.len());
220    let mut keep_by_class: HashMap<i64, Vec<usize>> = HashMap::new();
221    'candidate: for index in order {
222        let class_id = detections[index].class_id;
223        let selected_for_class = keep_by_class.entry(class_id).or_default();
224        for &selected in selected_for_class.iter() {
225            if iou_with_areas(
226                detections[index].bbox,
227                areas[index],
228                detections[selected].bbox,
229                areas[selected],
230            ) > iou_threshold
231            {
232                continue 'candidate;
233            }
234        }
235        keep.push(index);
236        selected_for_class.push(index);
237    }
238    Ok(keep)
239}
240
241/// Soft non-maximum suppression with deterministic score ordering.
242pub fn soft_nms(
243    boxes: &[BoundingBox2],
244    scores: &[f32],
245    score_threshold: f32,
246    iou_threshold: f32,
247    method: SoftNmsMethod,
248) -> VisionResult<Vec<ScoredIndex>> {
249    validate_nms_inputs(boxes, scores, score_threshold, iou_threshold)?;
250    if matches!(method, SoftNmsMethod::Gaussian { sigma } if !sigma.is_finite() || sigma <= 0.0) {
251        return Err(VisionError::InvalidParameter(
252            "Soft-NMS Gaussian sigma must be finite and positive".to_owned(),
253        ));
254    }
255    let mut candidates: Vec<ScoredIndex> = scores
256        .iter()
257        .copied()
258        .enumerate()
259        .map(|(index, score)| ScoredIndex { index, score })
260        .collect();
261    let areas: Vec<f32> = boxes.iter().copied().map(BoundingBox2::area).collect();
262    let mut output = Vec::with_capacity(candidates.len());
263    let mut best = best_scored_candidate(&candidates);
264    while let Some(best_index) = best {
265        if candidates[best_index].score < score_threshold {
266            break;
267        }
268        let selected = candidates.swap_remove(best_index);
269        output.push(selected);
270        let selected_box = boxes[selected.index];
271        let selected_area = areas[selected.index];
272        let mut active_count = 0;
273        let mut next_best = None;
274        for read_index in 0..candidates.len() {
275            let mut candidate = candidates[read_index];
276            let overlap = iou_with_areas(
277                selected_box,
278                selected_area,
279                boxes[candidate.index],
280                areas[candidate.index],
281            );
282            if overlap != 0.0 {
283                match method {
284                    SoftNmsMethod::Hard => {
285                        if overlap > iou_threshold {
286                            candidate.score *= 0.0;
287                        }
288                    }
289                    SoftNmsMethod::Linear => {
290                        if overlap > iou_threshold {
291                            candidate.score *= 1.0 - overlap;
292                        }
293                    }
294                    SoftNmsMethod::Gaussian { sigma } => {
295                        candidate.score *= (-(overlap * overlap) / sigma).exp();
296                    }
297                }
298            }
299            if candidate.score >= score_threshold {
300                candidates[active_count] = candidate;
301                let is_next_best = match next_best {
302                    Some(index) => scored_candidate_precedes(candidate, candidates[index]),
303                    None => true,
304                };
305                if is_next_best {
306                    next_best = Some(active_count);
307                }
308                active_count += 1;
309            }
310        }
311        candidates.truncate(active_count);
312        best = next_best;
313    }
314    Ok(output)
315}
316
317fn best_scored_candidate(candidates: &[ScoredIndex]) -> Option<usize> {
318    if candidates.is_empty() {
319        return None;
320    }
321    let mut best = 0;
322    for index in 1..candidates.len() {
323        if scored_candidate_precedes(candidates[index], candidates[best]) {
324            best = index;
325        }
326    }
327    Some(best)
328}
329
330fn scored_candidate_precedes(left: ScoredIndex, right: ScoredIndex) -> bool {
331    left.score > right.score || left.score == right.score && left.index < right.index
332}
333
334fn iou_with_areas(left: BoundingBox2, left_area: f32, right: BoundingBox2, right_area: f32) -> f32 {
335    let intersection_width = left.x_max.min(right.x_max) - left.x_min.max(right.x_min);
336    if intersection_width <= 0.0 {
337        return 0.0;
338    }
339    let intersection_height = left.y_max.min(right.y_max) - left.y_min.max(right.y_min);
340    if intersection_height <= 0.0 {
341        return 0.0;
342    }
343    let intersection = intersection_width * intersection_height;
344    let union = left_area + right_area - intersection;
345    if union > 0.0 {
346        intersection / union
347    } else {
348        0.0
349    }
350}
351
352fn validate_nms_inputs(
353    boxes: &[BoundingBox2],
354    scores: &[f32],
355    score_threshold: f32,
356    iou_threshold: f32,
357) -> VisionResult<()> {
358    if boxes.len() != scores.len() {
359        return Err(VisionError::ShapeMismatch(
360            "boxes and scores must have equal lengths".to_owned(),
361        ));
362    }
363    if !score_threshold.is_finite()
364        || !iou_threshold.is_finite()
365        || !(0.0..=1.0).contains(&iou_threshold)
366    {
367        return Err(VisionError::InvalidParameter(
368            "NMS thresholds must be finite and IoU in [0, 1]".to_owned(),
369        ));
370    }
371    if boxes.iter().any(|bbox| !bbox.is_valid()) || scores.iter().any(|score| !score.is_finite()) {
372        return Err(VisionError::InvalidParameter(
373            "boxes must be valid and scores finite".to_owned(),
374        ));
375    }
376    Ok(())
377}
378
379fn sort_indices_by_score(indices: &mut [usize], scores: &[f32]) {
380    indices.sort_by(|&left, &right| {
381        scores[right]
382            .partial_cmp(&scores[left])
383            .unwrap_or(Ordering::Equal)
384            .then_with(|| left.cmp(&right))
385    });
386}
387
388#[cfg(test)]
389mod tests {
390    use super::{
391        batched_nms, iou_with_areas, nms, soft_nms, BoundingBox2, Detection, SoftNmsMethod,
392    };
393
394    fn bbox(x0: f32, y0: f32, x1: f32, y1: f32) -> BoundingBox2 {
395        BoundingBox2::try_new(x0, y0, x1, y1).unwrap()
396    }
397
398    #[test]
399    fn iou_matches_known_overlap() {
400        let a = bbox(0.0, 0.0, 2.0, 2.0);
401        let b = bbox(1.0, 1.0, 3.0, 3.0);
402        assert!((a.iou(b) - 1.0 / 7.0).abs() < 1e-6);
403        assert!(a.generalized_iou(b) < a.iou(b));
404    }
405
406    #[test]
407    fn nms_suppresses_lower_scored_overlap() {
408        let boxes = [bbox(0.0, 0.0, 2.0, 2.0), bbox(0.1, 0.1, 2.1, 2.1), bbox(5.0, 5.0, 6.0, 6.0)];
409        assert_eq!(nms(&boxes, &[0.9, 0.8, 0.7], 0.0, 0.5).unwrap(), vec![0, 2]);
410    }
411
412    #[test]
413    fn cached_area_iou_matches_public_iou() {
414        let boxes = [
415            BoundingBox2::try_new(0.0, 0.0, 20.0, 20.0).unwrap(),
416            BoundingBox2::try_new(2.0, 2.0, 18.0, 18.0).unwrap(),
417            BoundingBox2::try_new(30.0, 30.0, 45.0, 45.0).unwrap(),
418            BoundingBox2::try_new(5.0, 5.0, 5.0, 8.0).unwrap(),
419        ];
420        for &left in &boxes {
421            for &right in &boxes {
422                assert_eq!(iou_with_areas(left, left.area(), right, right.area()), left.iou(right),);
423            }
424        }
425    }
426
427    #[test]
428    fn batched_nms_keeps_overlapping_different_classes() {
429        let detections = [
430            Detection { bbox: bbox(0.0, 0.0, 2.0, 2.0), score: 0.9, class_id: 1 },
431            Detection { bbox: bbox(0.0, 0.0, 2.0, 2.0), score: 0.8, class_id: 2 },
432        ];
433        assert_eq!(batched_nms(&detections, 0.0, 0.5).unwrap(), vec![0, 1]);
434    }
435
436    #[test]
437    fn batched_nms_suppresses_per_class_in_global_score_order() {
438        let detections = [
439            Detection { bbox: bbox(0.0, 0.0, 2.0, 2.0), score: 0.7, class_id: 4 },
440            Detection { bbox: bbox(0.1, 0.1, 2.1, 2.1), score: 0.9, class_id: 4 },
441            Detection { bbox: bbox(0.1, 0.1, 2.1, 2.1), score: 0.8, class_id: 9 },
442            Detection { bbox: bbox(5.0, 5.0, 6.0, 6.0), score: 0.6, class_id: 4 },
443        ];
444        assert_eq!(batched_nms(&detections, 0.0, 0.5).unwrap(), vec![1, 2, 3]);
445    }
446
447    #[test]
448    fn class_buckets_match_global_scan_reference() {
449        let mut state = 119_u64;
450        let detections = (0..257)
451            .map(|index| {
452                let x = sample(&mut state) * 64.0;
453                let y = sample(&mut state) * 64.0;
454                let width = 1.0 + sample(&mut state) * 20.0;
455                let height = 1.0 + sample(&mut state) * 20.0;
456                Detection {
457                    bbox: bbox(x, y, x + width, y + height),
458                    score: sample(&mut state),
459                    class_id: (index % 11) as i64 - 5,
460                }
461            })
462            .collect::<Vec<_>>();
463        let actual = batched_nms(&detections, 0.2, 0.45).unwrap();
464
465        let mut order = (0..detections.len())
466            .filter(|&index| detections[index].score >= 0.2)
467            .collect::<Vec<_>>();
468        order.sort_by(|&left, &right| {
469            detections[right]
470                .score
471                .partial_cmp(&detections[left].score)
472                .unwrap()
473                .then_with(|| left.cmp(&right))
474        });
475        let mut expected: Vec<usize> = Vec::new();
476        'candidate: for index in order {
477            for &selected in &expected {
478                if detections[index].class_id == detections[selected].class_id
479                    && detections[index].bbox.iou(detections[selected].bbox) > 0.45
480                {
481                    continue 'candidate;
482                }
483            }
484            expected.push(index);
485        }
486        assert_eq!(actual, expected);
487    }
488
489    fn sample(state: &mut u64) -> f32 {
490        *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
491        let mut value = *state;
492        value = (value ^ (value >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
493        value = (value ^ (value >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
494        value ^= value >> 31;
495        (value >> 40) as f32 / (1_u32 << 24) as f32
496    }
497
498    #[test]
499    fn soft_nms_decays_overlapping_score() {
500        let boxes = [bbox(0.0, 0.0, 2.0, 2.0), bbox(0.1, 0.1, 2.1, 2.1)];
501        let result = soft_nms(&boxes, &[0.9, 0.8], 0.01, 0.5, SoftNmsMethod::Linear).unwrap();
502        assert_eq!(result[0].index, 0);
503        assert!(result[1].score < 0.8);
504        assert!(soft_nms(&boxes, &[0.9, 0.8], 0.01, 0.5, SoftNmsMethod::Gaussian { sigma: 0.0 },)
505            .is_err());
506    }
507
508    #[test]
509    fn soft_nms_selection_scan_preserves_sorting_semantics() {
510        let mut state = 127_u64;
511        let boxes = (0..73)
512            .map(|_| {
513                let x = sample(&mut state) * 64.0;
514                let y = sample(&mut state) * 64.0;
515                let width = 1.0 + sample(&mut state) * 20.0;
516                let height = 1.0 + sample(&mut state) * 20.0;
517                bbox(x, y, x + width, y + height)
518            })
519            .collect::<Vec<_>>();
520        let mut scores = (0..boxes.len()).map(|_| sample(&mut state)).collect::<Vec<_>>();
521        scores[5] = scores[2];
522        for method in
523            [SoftNmsMethod::Hard, SoftNmsMethod::Linear, SoftNmsMethod::Gaussian { sigma: 0.5 }]
524        {
525            let actual = soft_nms(&boxes, &scores, 0.2, 0.45, method).unwrap();
526            let expected = sorting_soft_nms_reference(&boxes, &scores, 0.2, 0.45, method);
527            assert_eq!(actual, expected);
528        }
529
530        let overlapping = [bbox(0.0, 0.0, 2.0, 2.0), bbox(0.0, 0.0, 2.0, 2.0)];
531        let negative_scores = [1.0, -0.6];
532        assert_eq!(
533            soft_nms(&overlapping, &negative_scores, -0.5, 0.45, SoftNmsMethod::Hard).unwrap(),
534            sorting_soft_nms_reference(
535                &overlapping,
536                &negative_scores,
537                -0.5,
538                0.45,
539                SoftNmsMethod::Hard,
540            ),
541        );
542    }
543
544    fn sorting_soft_nms_reference(
545        boxes: &[BoundingBox2],
546        scores: &[f32],
547        score_threshold: f32,
548        iou_threshold: f32,
549        method: SoftNmsMethod,
550    ) -> Vec<super::ScoredIndex> {
551        let mut candidates = scores
552            .iter()
553            .copied()
554            .enumerate()
555            .map(|(index, score)| super::ScoredIndex { index, score })
556            .collect::<Vec<_>>();
557        let mut output = Vec::new();
558        while !candidates.is_empty() {
559            candidates.sort_by(|left, right| {
560                right
561                    .score
562                    .partial_cmp(&left.score)
563                    .unwrap()
564                    .then_with(|| left.index.cmp(&right.index))
565            });
566            let selected = candidates.remove(0);
567            if selected.score < score_threshold {
568                break;
569            }
570            output.push(selected);
571            for candidate in &mut candidates {
572                let overlap = boxes[selected.index].iou(boxes[candidate.index]);
573                let weight = match method {
574                    SoftNmsMethod::Hard => (overlap <= iou_threshold) as u8 as f32,
575                    SoftNmsMethod::Linear => {
576                        if overlap > iou_threshold {
577                            1.0 - overlap
578                        } else {
579                            1.0
580                        }
581                    }
582                    SoftNmsMethod::Gaussian { sigma } => (-(overlap * overlap) / sigma).exp(),
583                };
584                candidate.score *= weight;
585            }
586            candidates.retain(|candidate| candidate.score >= score_threshold);
587        }
588        output
589    }
590}
spatialrust_vision/detection.rs

spatialrust_vision/
detection.rs
