Skip to main content

edgefirst_decoder/
float.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Float-domain NMS and box geometry.
5//!
6//! Two groups of primitives, both reusable outside the decoder:
7//!
8//! - **Suppression** — [`nms_float`], [`nms_class_aware_float`], and the
9//!   `_extra` variants that carry a per-box payload (mask coefficients, say)
10//!   through suppression so the survivors stay paired with their data.
11//! - **Geometry** — [`intersection_area`], [`box_area`], [`iou_value`], and
12//!   [`ios_value`]. The tiled-inference merge in [`crate::tiling`] shares
13//!   these, which is what keeps its IoS metric numerically identical to the
14//!   IoU used by ordinary NMS.
15//!
16//! [`jaccard`] and [`jaccard_batch4`] are the threshold-test forms used on the
17//! suppression hot path; they answer "do these overlap past `iou`?" without
18//! materializing the ratio.
19
20use crate::{arg_max, BBoxTypeTrait, BoundingBox, DetectBox};
21use ndarray::{ArrayView1, ArrayView2};
22use num_traits::{AsPrimitive, Float};
23use rayon::iter::{IntoParallelIterator, ParallelIterator};
24use rayon::slice::ParallelSliceMut;
25
26/// Candidate filters over fewer score elements than this run on the calling
27/// thread: handing a few hundred rows to rayon costs more CPU, in waking and
28/// then spinning workers, than filtering them.
29const PAR_MIN_SCORES: usize = 1 << 16;
30
31/// `f` over every `(index, score row, box row)`, keeping the `Some`s in row
32/// order, on the rayon pool when the score array is large.
33fn filter_rows<BOX, SCORE, T, F>(boxes: ArrayView2<BOX>, scores: ArrayView2<SCORE>, f: F) -> Vec<T>
34where
35    BOX: Sync,
36    SCORE: Sync,
37    T: Send,
38    F: Fn(usize, ArrayView1<SCORE>, ArrayView1<BOX>) -> Option<T> + Sync + Send,
39{
40    let (n, nc) = scores.dim();
41    let row = |i: usize| f(i, scores.row(i), boxes.row(i));
42    if n * nc < PAR_MIN_SCORES {
43        (0..n).filter_map(row).collect()
44    } else {
45        (0..n).into_par_iter().filter_map(row).collect()
46    }
47}
48
49/// Post processes boxes and scores tensors into detection boxes, filtering out
50/// any boxes below the score threshold. The boxes tensor is converted to XYXY
51/// using the given BBoxTypeTrait. The order of the boxes is preserved.
52pub fn postprocess_boxes_float<
53    B: BBoxTypeTrait,
54    BOX: Float + AsPrimitive<f32> + Send + Sync,
55    SCORE: Float + AsPrimitive<f32> + Send + Sync,
56>(
57    threshold: SCORE,
58    boxes: ArrayView2<BOX>,
59    scores: ArrayView2<SCORE>,
60) -> Vec<DetectBox> {
61    assert_eq!(scores.dim().0, boxes.dim().0);
62    assert_eq!(boxes.dim().1, 4);
63    filter_rows(boxes, scores, |_, score, bbox| {
64        let (score_, label) = arg_max(score);
65        if score_ < threshold {
66            return None;
67        }
68
69        let bbox = B::ndarray_to_xyxy_float(bbox);
70        Some(DetectBox {
71            label,
72            score: score_.as_(),
73            bbox: bbox.into(),
74        })
75    })
76}
77
78/// Post processes boxes and scores tensors into detection boxes, filtering out
79/// any boxes below the score threshold. The boxes tensor is converted to XYXY
80/// using the given BBoxTypeTrait. The order of the boxes is preserved.
81///
82/// This function is very similar to `postprocess_boxes_float` but will also
83/// return the index of the box. The boxes will be in ascending index order.
84pub fn postprocess_boxes_index_float<
85    B: BBoxTypeTrait,
86    BOX: Float + AsPrimitive<f32> + Send + Sync,
87    SCORE: Float + AsPrimitive<f32> + Send + Sync,
88>(
89    threshold: SCORE,
90    boxes: ArrayView2<BOX>,
91    scores: ArrayView2<SCORE>,
92) -> Vec<(DetectBox, usize)> {
93    assert_eq!(scores.dim().0, boxes.dim().0);
94    assert_eq!(boxes.dim().1, 4);
95    filter_rows(boxes, scores, |i, score, bbox| {
96        let (score_, label) = arg_max(score);
97        if score_ < threshold {
98            return None;
99        }
100
101        let bbox = B::ndarray_to_xyxy_float(bbox);
102        Some((
103            DetectBox {
104                label,
105                score: score_.as_(),
106                bbox: bbox.into(),
107            },
108            i,
109        ))
110    })
111}
112
113/// Multi-label variant of [`postprocess_boxes_index_float`].
114///
115/// For each anchor row, emits one `(DetectBox, anchor_idx)` per class whose
116/// score meets `threshold` — every class, not just the argmax.  The same
117/// `anchor_idx` is returned for all per-class entries of a given anchor so
118/// that downstream mask-coefficient lookup can reuse the shared coefficient
119/// row.
120///
121/// The bbox is computed once per anchor (via `B::ndarray_to_xyxy_float`) and
122/// reused across all emitted classes, avoiding redundant work.
123///
124/// Intended for **validation/mAP evaluation only** (the Ultralytics `val`
125/// convention).  Deployment must use the argmax variant
126/// [`postprocess_boxes_index_float`] so trackers see at most one box per
127/// anchor.
128pub fn postprocess_boxes_multilabel_index_float<
129    B: BBoxTypeTrait,
130    BOX: Float + AsPrimitive<f32> + Send + Sync,
131    SCORE: Float + AsPrimitive<f32> + Send + Sync,
132>(
133    threshold: SCORE,
134    boxes: ArrayView2<BOX>,
135    scores: ArrayView2<SCORE>,
136) -> Vec<(DetectBox, usize)> {
137    assert_eq!(scores.dim().0, boxes.dim().0);
138    assert_eq!(boxes.dim().1, 4);
139    filter_rows(boxes, scores, |anchor_idx, score_row, bbox_row| {
140        if !score_row.iter().any(|&s| s >= threshold) {
141            return None;
142        }
143        // Compute bbox once; clone into each per-class candidate.
144        let bbox = B::ndarray_to_xyxy_float(bbox_row);
145        let bbox: crate::BoundingBox = bbox.into();
146        Some(
147            score_row
148                .iter()
149                .enumerate()
150                .filter(|(_, &s)| s >= threshold)
151                .map(|(c, &s)| {
152                    (
153                        DetectBox {
154                            label: c,
155                            score: s.as_(),
156                            bbox,
157                        },
158                        anchor_idx,
159                    )
160                })
161                .collect::<Vec<_>>(),
162        )
163    })
164    .into_iter()
165    .flatten()
166    .collect()
167}
168
169/// [`postprocess_boxes_multilabel_index_float`] without anchor indices.
170///
171/// # Errors
172/// [`crate::DecoderError::InvalidShape`] unless boxes are `[N, 4]` and scores
173/// `[N, C]`.
174pub fn postprocess_boxes_multilabel_float<
175    B: BBoxTypeTrait,
176    BOX: Float + AsPrimitive<f32> + Send + Sync,
177    SCORE: Float + AsPrimitive<f32> + Send + Sync,
178>(
179    threshold: SCORE,
180    boxes: ArrayView2<BOX>,
181    scores: ArrayView2<SCORE>,
182) -> crate::DecoderResult<Vec<DetectBox>> {
183    crate::byte::check_candidate_shapes(boxes.dim(), scores.dim())?;
184    Ok(
185        postprocess_boxes_multilabel_index_float::<B, _, _>(threshold, boxes, scores)
186            .into_iter()
187            .map(|(b, _)| b)
188            .collect(),
189    )
190}
191
192/// Uses NMS to filter boxes based on the score and iou. Sorts boxes by score,
193/// then greedily selects a subset of boxes in descending order of score.
194///
195/// If `max_det` is `Some(n)`, the greedy loop stops as soon as `n` survivors
196/// have been confirmed. Because the input is sorted descending, the first `n`
197/// survivors are the highest-scoring `n`, so the post-NMS top-`n` is preserved
198/// without iterating the full O(N²) suppression loop.
199#[must_use]
200pub fn nms_float(iou: f32, max_det: Option<usize>, mut boxes: Vec<DetectBox>) -> Vec<DetectBox> {
201    // Boxes get sorted by score in descending order so we know based on the
202    // index the scoring of the boxes and can skip parts of the loop.
203    boxes.par_sort_by(|a, b| b.score.total_cmp(&a.score));
204
205    // When the iou is 1.0 or larger, no boxes will be filtered so we just return
206    // immediately
207    if iou >= 1.0 {
208        return match max_det {
209            Some(n) => {
210                boxes.truncate(n);
211                boxes
212            }
213            None => boxes,
214        };
215    }
216
217    let cap = max_det.unwrap_or(usize::MAX);
218    let mut survivors: usize = 0;
219
220    // Outer loop over all boxes.
221    for i in 0..boxes.len() {
222        if boxes[i].score < 0.0 {
223            // this box was merged with a different box earlier
224            continue;
225        }
226        for j in (i + 1)..boxes.len() {
227            // Inner loop over boxes with lower score (later in the list).
228
229            if boxes[j].score < 0.0 {
230                // this box was suppressed by different box earlier
231                continue;
232            }
233            if jaccard(&boxes[j].bbox, &boxes[i].bbox, iou) {
234                // max_box(boxes[j].bbox, &mut boxes[i].bbox);
235                boxes[j].score = -1.0;
236            }
237        }
238
239        // NOTE: jaccard_batch4_neon is available for callers that can
240        // batch unsuppressed candidates externally. It is not used
241        // inline here because score=-1 marking creates sparse gaps
242        // that prevent contiguous 4-box batching.
243        survivors += 1;
244        if survivors >= cap {
245            break;
246        }
247    }
248    // Filter out suppressed boxes; cap at `max_det` because boxes after the
249    // break may still hold positive scores but score lower than every survivor.
250    boxes
251        .into_iter()
252        .filter(|b| b.score >= 0.0)
253        .take(cap)
254        .collect()
255}
256
257/// Uses NMS to filter boxes based on the score and iou. Sorts boxes by score,
258/// then greedily selects a subset of boxes in descending order of score.
259///
260/// This is same as `nms_float` but will also include extra information along
261/// with each box, such as the index
262#[must_use]
263pub fn nms_extra_float<E: Send + Sync>(
264    iou: f32,
265    max_det: Option<usize>,
266    mut boxes: Vec<(DetectBox, E)>,
267) -> Vec<(DetectBox, E)> {
268    // Boxes get sorted by score in descending order so we know based on the
269    // index the scoring of the boxes and can skip parts of the loop.
270    boxes.par_sort_by(|a, b| b.0.score.total_cmp(&a.0.score));
271
272    // When the iou is 1.0 or larger, no boxes will be filtered so we just return
273    // immediately
274    if iou >= 1.0 {
275        return match max_det {
276            Some(n) => {
277                boxes.truncate(n);
278                boxes
279            }
280            None => boxes,
281        };
282    }
283
284    let cap = max_det.unwrap_or(usize::MAX);
285    let mut survivors: usize = 0;
286
287    // Outer loop over all boxes.
288    for i in 0..boxes.len() {
289        if boxes[i].0.score < 0.0 {
290            // this box was merged with a different box earlier
291            continue;
292        }
293        for j in (i + 1)..boxes.len() {
294            // Inner loop over boxes with lower score (later in the list).
295
296            if boxes[j].0.score < 0.0 {
297                // this box was suppressed by different box earlier
298                continue;
299            }
300            if jaccard(&boxes[j].0.bbox, &boxes[i].0.bbox, iou) {
301                // max_box(boxes[j].bbox, &mut boxes[i].bbox);
302                boxes[j].0.score = -1.0;
303            }
304        }
305        survivors += 1;
306        if survivors >= cap {
307            break;
308        }
309    }
310
311    // Filter out suppressed boxes; cap at `max_det` for the same reason as
312    // `nms_float`.
313    boxes
314        .into_iter()
315        .filter(|b| b.0.score >= 0.0)
316        .take(cap)
317        .collect()
318}
319
320/// Class-aware NMS: only suppress boxes with the same label.
321///
322/// Sorts boxes by score, then greedily selects a subset of boxes in descending
323/// order of score. Unlike class-agnostic NMS, boxes are only suppressed if they
324/// have the same class label AND overlap above the IoU threshold.
325///
326/// # Example
327/// ```
328/// # use edgefirst_decoder::{BoundingBox, DetectBox, float::nms_class_aware_float};
329/// let boxes = vec![
330///     DetectBox {
331///         bbox: BoundingBox::new(0.0, 0.0, 0.5, 0.5),
332///         score: 0.9,
333///         label: 0,
334///     },
335///     DetectBox {
336///         bbox: BoundingBox::new(0.1, 0.1, 0.6, 0.6),
337///         score: 0.8,
338///         label: 1,
339///     }, // different class
340/// ];
341/// // Both boxes survive because they have different labels
342/// let result = nms_class_aware_float(0.3, None, boxes);
343/// assert_eq!(result.len(), 2);
344/// ```
345#[must_use]
346pub fn nms_class_aware_float(
347    iou: f32,
348    max_det: Option<usize>,
349    mut boxes: Vec<DetectBox>,
350) -> Vec<DetectBox> {
351    boxes.par_sort_by(|a, b| b.score.total_cmp(&a.score));
352
353    if iou >= 1.0 {
354        return match max_det {
355            Some(n) => {
356                boxes.truncate(n);
357                boxes
358            }
359            None => boxes,
360        };
361    }
362
363    let cap = max_det.unwrap_or(usize::MAX);
364    let mut survivors: usize = 0;
365
366    for i in 0..boxes.len() {
367        if boxes[i].score < 0.0 {
368            continue;
369        }
370        for j in (i + 1)..boxes.len() {
371            if boxes[j].score < 0.0 {
372                continue;
373            }
374            // Only suppress if same class AND overlapping
375            if boxes[j].label == boxes[i].label && jaccard(&boxes[j].bbox, &boxes[i].bbox, iou) {
376                boxes[j].score = -1.0;
377            }
378        }
379        survivors += 1;
380        if survivors >= cap {
381            break;
382        }
383    }
384    boxes
385        .into_iter()
386        .filter(|b| b.score >= 0.0)
387        .take(cap)
388        .collect()
389}
390
391/// Class-aware NMS with extra data: only suppress boxes with the same label.
392///
393/// This is same as `nms_class_aware_float` but will also include extra
394/// information along with each box, such as the index.
395#[must_use]
396pub fn nms_extra_class_aware_float<E: Send + Sync>(
397    iou: f32,
398    max_det: Option<usize>,
399    mut boxes: Vec<(DetectBox, E)>,
400) -> Vec<(DetectBox, E)> {
401    boxes.par_sort_by(|a, b| b.0.score.total_cmp(&a.0.score));
402
403    // When the iou is 1.0 or larger, no boxes will be filtered so we just return
404    // immediately
405    if iou >= 1.0 {
406        return match max_det {
407            Some(n) => {
408                boxes.truncate(n);
409                boxes
410            }
411            None => boxes,
412        };
413    }
414
415    let cap = max_det.unwrap_or(usize::MAX);
416    let mut survivors: usize = 0;
417
418    for i in 0..boxes.len() {
419        if boxes[i].0.score < 0.0 {
420            continue;
421        }
422        for j in (i + 1)..boxes.len() {
423            if boxes[j].0.score < 0.0 {
424                continue;
425            }
426            // Only suppress if same class AND overlapping
427            if boxes[j].0.label == boxes[i].0.label
428                && jaccard(&boxes[j].0.bbox, &boxes[i].0.bbox, iou)
429            {
430                boxes[j].0.score = -1.0;
431            }
432        }
433        survivors += 1;
434        if survivors >= cap {
435            break;
436        }
437    }
438    boxes
439        .into_iter()
440        .filter(|b| b.0.score >= 0.0)
441        .take(cap)
442        .collect()
443}
444
445/// Area of the axis-aligned intersection of two boxes, clamped to `>= 0` per
446/// side. Mirrors the `inter` term in ModelPack `metrics/tiled.py::_match`.
447#[must_use]
448#[inline]
449pub fn intersection_area(a: &BoundingBox, b: &BoundingBox) -> f32 {
450    let left = a.xmin.max(b.xmin);
451    let top = a.ymin.max(b.ymin);
452    let right = a.xmax.min(b.xmax);
453    let bottom = a.ymax.min(b.ymax);
454    (right - left).max(0.0) * (bottom - top).max(0.0)
455}
456
457/// Area of a box, clamped to `>= 0` per side. Mirrors ModelPack `_areas`.
458#[must_use]
459#[inline]
460pub fn box_area(b: &BoundingBox) -> f32 {
461    (b.xmax - b.xmin).max(0.0) * (b.ymax - b.ymin).max(0.0)
462}
463
464/// Intersection-over-Union value in `[0, 1]`. Mirrors ModelPack
465/// `_match(metric='iou')`: `inter / max(area_a + area_b - inter, 1e-9)`.
466#[must_use]
467#[inline]
468pub fn iou_value(a: &BoundingBox, b: &BoundingBox) -> f32 {
469    let inter = intersection_area(a, b);
470    let union = (box_area(a) + box_area(b) - inter).max(1e-9);
471    inter / union
472}
473
474/// Intersection-over-Smaller value in `[0, 1]`. Mirrors ModelPack
475/// `_match(metric='ios')`: `inter / max(min(area_a, area_b), 1e-9)`. Used by the
476/// tiled-detection merge: a seam-split object has low IoU but high IoS.
477#[must_use]
478#[inline]
479pub fn ios_value(a: &BoundingBox, b: &BoundingBox) -> f32 {
480    let inter = intersection_area(a, b);
481    let denom = box_area(a).min(box_area(b)).max(1e-9);
482    inter / denom
483}
484
485/// Returns true if the IOU of the given bounding boxes is greater than the iou
486/// threshold
487///
488/// Kept as a standalone inline (NOT routed through [`iou_value`]) so this hot
489/// NMS primitive stays byte-identical to its NEON sibling
490/// [`jaccard_batch4`]; the value helpers above are additive and used by the
491/// tiled-detection merge.
492///
493/// # Example
494/// ```
495/// # use edgefirst_decoder::{BoundingBox, float::jaccard};
496/// let a = BoundingBox::new(0.0, 0.0, 0.2, 0.2);
497/// let b = BoundingBox::new(0.1, 0.1, 0.3, 0.3);
498/// let iou_threshold = 0.1;
499/// let result = jaccard(&a, &b, iou_threshold);
500/// assert!(result);
501/// ```
502pub fn jaccard(a: &BoundingBox, b: &BoundingBox, iou: f32) -> bool {
503    let left = a.xmin.max(b.xmin);
504    let top = a.ymin.max(b.ymin);
505    let right = a.xmax.min(b.xmax);
506    let bottom = a.ymax.min(b.ymax);
507
508    let intersection = (right - left).max(0.0) * (bottom - top).max(0.0);
509    let area_a = (a.xmax - a.xmin) * (a.ymax - a.ymin);
510    let area_b = (b.xmax - b.xmin) * (b.ymax - b.ymin);
511
512    // need to make sure we are not dividing by zero
513    let union = area_a + area_b - intersection;
514
515    intersection > iou * union
516}
517
518/// Batch IoU check: test one reference box `a` against 4 candidate boxes.
519///
520/// Returns a 4-element array of booleans: `result[i]` is true if
521/// `jaccard(a, boxes[i], iou)` would return true.
522///
523/// On aarch64, uses NEON `vmaxq_f32`/`vminq_f32` for vectorized
524/// intersection computation. On other architectures falls back to
525/// 4 scalar `jaccard` calls.
526#[inline]
527pub fn jaccard_batch4(a: &BoundingBox, boxes: &[BoundingBox; 4], iou: f32) -> [bool; 4] {
528    #[cfg(target_arch = "aarch64")]
529    {
530        // SAFETY: NEON is mandatory on aarch64.
531        unsafe { jaccard_batch4_neon(a, boxes, iou) }
532    }
533    #[cfg(not(target_arch = "aarch64"))]
534    {
535        [
536            jaccard(a, &boxes[0], iou),
537            jaccard(a, &boxes[1], iou),
538            jaccard(a, &boxes[2], iou),
539            jaccard(a, &boxes[3], iou),
540        ]
541    }
542}
543
544/// NEON-vectorized batch IoU for 4 candidate boxes against one reference.
545///
546/// Loads xmin/ymin/xmax/ymax of the 4 candidates into separate NEON
547/// registers (AoS→SoA transpose), then computes intersection, union,
548/// and the `intersection > iou * union` test in 4-wide SIMD.
549#[cfg(target_arch = "aarch64")]
550#[target_feature(enable = "neon")]
551unsafe fn jaccard_batch4_neon(a: &BoundingBox, boxes: &[BoundingBox; 4], iou: f32) -> [bool; 4] {
552    unsafe {
553        use std::arch::aarch64::*; // NOSONAR
554
555        let zero = vdupq_n_f32(0.0);
556        let iou_v = vdupq_n_f32(iou);
557
558        // Reference box broadcast.
559        let a_xmin = vdupq_n_f32(a.xmin);
560        let a_ymin = vdupq_n_f32(a.ymin);
561        let a_xmax = vdupq_n_f32(a.xmax);
562        let a_ymax = vdupq_n_f32(a.ymax);
563        let area_a = vmulq_f32(vsubq_f32(a_xmax, a_xmin), vsubq_f32(a_ymax, a_ymin));
564
565        // Load 4 boxes (each BoundingBox is [xmin, ymin, xmax, ymax]).
566        let b0 = vld1q_f32(&boxes[0].xmin as *const f32);
567        let b1 = vld1q_f32(&boxes[1].xmin as *const f32);
568        let b2 = vld1q_f32(&boxes[2].xmin as *const f32);
569        let b3 = vld1q_f32(&boxes[3].xmin as *const f32);
570
571        // AoS → SoA transpose (4×4).
572        let t01_lo = vtrn1q_f32(b0, b1); // xmin0,xmin1,xmax0,xmax1
573        let t01_hi = vtrn2q_f32(b0, b1); // ymin0,ymin1,ymax0,ymax1
574        let t23_lo = vtrn1q_f32(b2, b3);
575        let t23_hi = vtrn2q_f32(b2, b3);
576
577        let b_xmin = vreinterpretq_f32_f64(vtrn1q_f64(
578            vreinterpretq_f64_f32(t01_lo),
579            vreinterpretq_f64_f32(t23_lo),
580        ));
581        let b_ymin = vreinterpretq_f32_f64(vtrn1q_f64(
582            vreinterpretq_f64_f32(t01_hi),
583            vreinterpretq_f64_f32(t23_hi),
584        ));
585        let b_xmax = vreinterpretq_f32_f64(vtrn2q_f64(
586            vreinterpretq_f64_f32(t01_lo),
587            vreinterpretq_f64_f32(t23_lo),
588        ));
589        let b_ymax = vreinterpretq_f32_f64(vtrn2q_f64(
590            vreinterpretq_f64_f32(t01_hi),
591            vreinterpretq_f64_f32(t23_hi),
592        ));
593
594        // Intersection.
595        let left = vmaxq_f32(a_xmin, b_xmin);
596        let top = vmaxq_f32(a_ymin, b_ymin);
597        let right = vminq_f32(a_xmax, b_xmax);
598        let bottom = vminq_f32(a_ymax, b_ymax);
599        let w = vmaxq_f32(vsubq_f32(right, left), zero);
600        let h = vmaxq_f32(vsubq_f32(bottom, top), zero);
601        let intersection = vmulq_f32(w, h);
602
603        // Area B.
604        let area_b = vmulq_f32(vsubq_f32(b_xmax, b_xmin), vsubq_f32(b_ymax, b_ymin));
605
606        // Union = area_a + area_b - intersection.
607        let union = vsubq_f32(vaddq_f32(area_a, area_b), intersection);
608
609        // Test: intersection > iou * union (equivalent to IoU > threshold).
610        let iou_union = vmulq_f32(iou_v, union);
611        let mask = vcgtq_f32(intersection, iou_union);
612
613        // Extract per-lane results.
614        [
615            vgetq_lane_u32(mask, 0) != 0,
616            vgetq_lane_u32(mask, 1) != 0,
617            vgetq_lane_u32(mask, 2) != 0,
618            vgetq_lane_u32(mask, 3) != 0,
619        ]
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626    use crate::BoundingBox;
627
628    #[test]
629    fn multilabel_float_drops_indices_only() {
630        let boxes = ndarray::arr2(&[[0.1f32, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4]]);
631        let scores = ndarray::arr2(&[[0.6f32, 0.1, 0.7], [0.05, 0.9, 0.05]]);
632        let with_idx = postprocess_boxes_multilabel_index_float::<crate::XYXY, _, _>(
633            0.5,
634            boxes.view(),
635            scores.view(),
636        );
637        let without = postprocess_boxes_multilabel_float::<crate::XYXY, _, _>(
638            0.5,
639            boxes.view(),
640            scores.view(),
641        )
642        .unwrap();
643        let stripped: Vec<_> = with_idx.into_iter().map(|(b, _)| b).collect();
644        assert_eq!(stripped, without);
645        assert_eq!(without.len(), 3);
646    }
647
648    #[test]
649    fn multilabel_float_rejects_non_four_column_boxes() {
650        let boxes = ndarray::arr2(&[[0.1f32, 0.1, 0.2, 0.2, 0.0]]);
651        let scores = ndarray::arr2(&[[0.6f32]]);
652        let r = postprocess_boxes_multilabel_float::<crate::XYXY, _, _>(
653            0.5,
654            boxes.view(),
655            scores.view(),
656        );
657        assert!(matches!(r, Err(crate::DecoderError::InvalidShape(_))));
658    }
659
660    /// Helper: create `n` non-overlapping boxes with descending f32 scores.
661    fn make_nms_boxes_float(n: usize) -> Vec<DetectBox> {
662        (0..n)
663            .map(|i| DetectBox {
664                bbox: BoundingBox {
665                    xmin: i as f32 * 100.0,
666                    ymin: 0.0,
667                    xmax: i as f32 * 100.0 + 10.0,
668                    ymax: 10.0,
669                },
670                label: 0,
671                score: 1.0 - i as f32 * 0.01,
672            })
673            .collect()
674    }
675
676    /// Two boxes at the same place with the same score. Sorting is stable
677    /// enough that the first stays and the second is suppressed; a score of
678    /// exactly 0.0 must not be mistaken for the `-1.0` suppression marker at
679    /// either end of the comparison.
680    fn overlapping_pair(score: f32) -> Vec<DetectBox> {
681        vec![
682            DetectBox {
683                bbox: BoundingBox {
684                    xmin: 0.0,
685                    ymin: 0.0,
686                    xmax: 10.0,
687                    ymax: 10.0,
688                },
689                label: 0,
690                score,
691            },
692            DetectBox {
693                bbox: BoundingBox {
694                    xmin: 1.0,
695                    ymin: 1.0,
696                    xmax: 11.0,
697                    ymax: 11.0,
698                },
699                label: 0,
700                score,
701            },
702        ]
703    }
704
705    /// A box scoring exactly 0.0 is a real detection that passed the
706    /// threshold, not a suppressed one. Treating `< 0.0` as `<= 0.0` would
707    /// skip it and leave its overlapping neighbour unsuppressed.
708    #[test]
709    fn nms_float_suppresses_overlapping_zero_score_boxes() {
710        let result = nms_float(0.5, None, overlapping_pair(0.0));
711        assert_eq!(result.len(), 1, "{result:?}");
712        assert_eq!(result[0].score, 0.0);
713    }
714
715    /// Three boxes in a chain: A overlaps B, B overlaps C, A does not reach
716    /// C. Once B is suppressed it must take no further part, or it suppresses
717    /// C on behalf of a box that is not in the output.
718    ///
719    /// The offsets carry the whole test. At 10 wide and stepping 2, A-B and
720    /// B-C are both 8/12 = 0.667 and A-C is 6/14 = 0.429, so only the
721    /// adjacent pairs clear the 0.5 threshold. Step further apart -- 0, 1 and
722    /// 9.5, say -- and B-C falls to 0.081, B cannot reach C whatever the
723    /// implementation does, and all three tests below pass with the bug they
724    /// exist to catch.
725    fn suppression_chain() -> Vec<DetectBox> {
726        let mk = |xmin: f32, xmax: f32, score: f32| DetectBox {
727            bbox: BoundingBox {
728                xmin,
729                ymin: 0.0,
730                xmax,
731                ymax: 10.0,
732            },
733            label: 0,
734            score,
735        };
736        vec![mk(0.0, 10.0, 0.9), mk(2.0, 12.0, 0.8), mk(4.0, 14.0, 0.7)]
737    }
738
739    #[test]
740    fn nms_float_does_not_suppress_on_behalf_of_a_suppressed_box() {
741        let result = nms_float(0.5, None, suppression_chain());
742        assert_eq!(result.len(), 2, "{result:?}");
743        assert_eq!(result[0].score, 0.9);
744        assert_eq!(result[1].score, 0.7);
745    }
746
747    #[test]
748    fn nms_extra_float_suppresses_overlapping_zero_score_boxes() {
749        let boxes = overlapping_pair(0.0)
750            .into_iter()
751            .enumerate()
752            .map(|(i, b)| (b, i))
753            .collect();
754        let result = nms_extra_float(0.5, None, boxes);
755        assert_eq!(result.len(), 1, "{result:?}");
756        // Both boxes score 0.0 and par_sort_by is an unstable sort, so which
757        // of the two survives is not fixed; that it carries its own payload
758        // is.
759        assert!(result[0].1 == 0 || result[0].1 == 1, "{result:?}");
760    }
761
762    #[test]
763    fn nms_extra_float_does_not_suppress_on_behalf_of_a_suppressed_box() {
764        let boxes = suppression_chain()
765            .into_iter()
766            .enumerate()
767            .map(|(i, b)| (b, i))
768            .collect();
769        let result = nms_extra_float(0.5, None, boxes);
770        assert_eq!(result.len(), 2, "{result:?}");
771        assert_eq!(result[0].1, 0);
772        assert_eq!(result[1].1, 2);
773    }
774
775    /// Class-aware NMS with no cap still has to run the whole suppression
776    /// loop. Breaking out of it early leaves overlapping same-class boxes in
777    /// the output.
778    #[test]
779    fn nms_class_aware_float_suppresses_same_class_without_a_cap() {
780        let mut boxes = overlapping_pair(0.9);
781        boxes[1].score = 0.8;
782        let result = nms_class_aware_float(0.5, None, boxes);
783        assert_eq!(result.len(), 1, "{result:?}");
784        assert_eq!(result[0].score, 0.9);
785    }
786
787    /// The same two boxes with different labels are both kept: class-aware
788    /// NMS suppresses within a class, never across.
789    #[test]
790    fn nms_class_aware_float_keeps_overlapping_boxes_of_other_classes() {
791        let mut boxes = overlapping_pair(0.9);
792        boxes[1].score = 0.8;
793        boxes[1].label = 1;
794        let result = nms_class_aware_float(0.5, None, boxes);
795        assert_eq!(result.len(), 2, "{result:?}");
796    }
797
798    #[test]
799    fn nms_class_aware_float_suppresses_overlapping_zero_score_boxes() {
800        let result = nms_class_aware_float(0.5, None, overlapping_pair(0.0));
801        assert_eq!(result.len(), 1, "{result:?}");
802    }
803
804    #[test]
805    fn nms_class_aware_float_does_not_suppress_on_behalf_of_a_suppressed_box() {
806        let result = nms_class_aware_float(0.5, None, suppression_chain());
807        assert_eq!(result.len(), 2, "{result:?}");
808        assert_eq!(result[1].score, 0.7);
809    }
810
811    /// A box that overlaps nothing, then a pair that overlaps each other. The
812    /// first box confirms a survivor without suppressing anything, so a loop
813    /// that stops at the first survivor leaves the pair un-deduplicated while
814    /// the three-box chain above, where the first box does the suppressing,
815    /// looks identical either way.
816    fn isolated_then_overlapping_pair() -> Vec<DetectBox> {
817        let mk = |xmin: f32, xmax: f32, score: f32| DetectBox {
818            bbox: BoundingBox {
819                xmin,
820                ymin: 0.0,
821                xmax,
822                ymax: 10.0,
823            },
824            label: 0,
825            score,
826        };
827        vec![
828            mk(100.0, 110.0, 0.9),
829            mk(0.0, 10.0, 0.8),
830            mk(1.0, 11.0, 0.7),
831        ]
832    }
833
834    #[test]
835    fn nms_class_aware_float_keeps_suppressing_after_the_first_survivor() {
836        let result = nms_class_aware_float(0.5, None, isolated_then_overlapping_pair());
837        assert_eq!(result.len(), 2, "{result:?}");
838        assert_eq!(result[0].score, 0.9);
839        assert_eq!(result[1].score, 0.8);
840    }
841
842    fn with_index(boxes: Vec<DetectBox>) -> Vec<(DetectBox, usize)> {
843        boxes.into_iter().enumerate().map(|(i, b)| (b, i)).collect()
844    }
845
846    #[test]
847    fn nms_extra_class_aware_float_keeps_suppressing_after_the_first_survivor() {
848        let result =
849            nms_extra_class_aware_float(0.5, None, with_index(isolated_then_overlapping_pair()));
850        assert_eq!(result.len(), 2, "{result:?}");
851        assert_eq!(result[1].1, 1);
852    }
853
854    #[test]
855    fn nms_extra_class_aware_float_suppresses_same_class_without_a_cap() {
856        let mut boxes = overlapping_pair(0.9);
857        boxes[1].score = 0.8;
858        let result = nms_extra_class_aware_float(0.5, None, with_index(boxes));
859        assert_eq!(result.len(), 1, "{result:?}");
860        assert_eq!(result[0].0.score, 0.9);
861        assert_eq!(result[0].1, 0);
862    }
863
864    #[test]
865    fn nms_extra_class_aware_float_keeps_overlapping_boxes_of_other_classes() {
866        let mut boxes = overlapping_pair(0.9);
867        boxes[1].score = 0.8;
868        boxes[1].label = 1;
869        let result = nms_extra_class_aware_float(0.5, None, with_index(boxes));
870        assert_eq!(result.len(), 2, "{result:?}");
871    }
872
873    #[test]
874    fn nms_extra_class_aware_float_suppresses_overlapping_zero_score_boxes() {
875        let result = nms_extra_class_aware_float(0.5, None, with_index(overlapping_pair(0.0)));
876        assert_eq!(result.len(), 1, "{result:?}");
877    }
878
879    #[test]
880    fn nms_extra_class_aware_float_does_not_suppress_on_behalf_of_a_suppressed_box() {
881        let result = nms_extra_class_aware_float(0.5, None, with_index(suppression_chain()));
882        assert_eq!(result.len(), 2, "{result:?}");
883        assert_eq!(result[1].1, 2);
884    }
885
886    /// The threshold is a floor, not a cut: a score equal to it is kept.
887    #[test]
888    fn postprocess_boxes_float_keeps_a_score_equal_to_the_threshold() {
889        let boxes = ndarray::array![[0.0_f32, 0.0, 10.0, 10.0], [0.0, 0.0, 10.0, 10.0]];
890        let scores = ndarray::array![[0.5_f32], [0.25]];
891
892        let kept = postprocess_boxes_float::<crate::XYXY, _, _>(0.5, boxes.view(), scores.view());
893        assert_eq!(kept.len(), 1, "{kept:?}");
894        assert_eq!(kept[0].score, 0.5);
895
896        let kept =
897            postprocess_boxes_index_float::<crate::XYXY, _, _>(0.5, boxes.view(), scores.view());
898        assert_eq!(kept.len(), 1, "{kept:?}");
899        assert_eq!(kept[0].1, 0);
900    }
901
902    #[test]
903    fn nms_float_max_det_matches_full_truncated() {
904        let boxes = make_nms_boxes_float(20);
905        let n = 5;
906        let full = nms_float(0.5, None, boxes.clone());
907        let capped = nms_float(0.5, Some(n), boxes);
908        assert_eq!(capped.len(), n);
909        for (f, c) in full[..n].iter().zip(capped.iter()) {
910            assert_eq!(f.bbox, c.bbox);
911            assert_eq!(f.score, c.score);
912        }
913    }
914
915    #[test]
916    fn nms_float_max_det_zero_returns_empty() {
917        let boxes = make_nms_boxes_float(10);
918        let result = nms_float(0.5, Some(0), boxes);
919        assert!(result.is_empty());
920    }
921
922    #[test]
923    fn nms_float_max_det_iou_ge_1_returns_sorted_truncated() {
924        let boxes = make_nms_boxes_float(10);
925        let result = nms_float(1.0, Some(3), boxes);
926        assert_eq!(result.len(), 3);
927        assert!(result[0].score >= result[1].score);
928        assert!(result[1].score >= result[2].score);
929    }
930
931    #[test]
932    fn nms_float_max_det_larger_than_input() {
933        let boxes = make_nms_boxes_float(5);
934        let full = nms_float(0.5, None, boxes.clone());
935        let capped = nms_float(0.5, Some(100), boxes);
936        assert_eq!(full.len(), capped.len());
937    }
938
939    // Parity literals from ModelPack metrics/tiled.py::_match for
940    // A=[100,100,400,300], B=[350,100,400,300] (B fully inside A):
941    // inter=10000, area_a=60000, area_b=10000 => IoS=1.0, IoU=0.16667.
942    #[test]
943    fn metric_values_match_modelpack_match() {
944        let a = BoundingBox::new(100.0, 100.0, 400.0, 300.0);
945        let b = BoundingBox::new(350.0, 100.0, 400.0, 300.0);
946        assert!((intersection_area(&a, &b) - 10000.0).abs() < 1e-3);
947        assert!((box_area(&a) - 60000.0).abs() < 1e-3);
948        assert!((box_area(&b) - 10000.0).abs() < 1e-3);
949        assert!((ios_value(&a, &b) - 1.0).abs() < 1e-6);
950        assert!((iou_value(&a, &b) - (1.0 / 6.0)).abs() < 1e-6);
951    }
952
953    #[test]
954    fn metric_values_disjoint_are_zero() {
955        let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
956        let b = BoundingBox::new(20.0, 20.0, 30.0, 30.0);
957        assert_eq!(intersection_area(&a, &b), 0.0);
958        assert_eq!(ios_value(&a, &b), 0.0);
959        assert_eq!(iou_value(&a, &b), 0.0);
960    }
961
962    #[test]
963    fn metric_values_zero_area_box_no_panic() {
964        // Degenerate (zero-area) box: denom floored to 1e-9, result finite.
965        let degenerate = BoundingBox::new(10.0, 10.0, 10.0, 20.0);
966        let other = BoundingBox::new(0.0, 0.0, 30.0, 30.0);
967        assert_eq!(box_area(&degenerate), 0.0);
968        assert!(ios_value(&degenerate, &other).is_finite());
969        assert!(iou_value(&degenerate, &other).is_finite());
970    }
971
972    #[test]
973    fn ios_high_where_iou_low_for_contained_box() {
974        // The SAHI seam case: a small box fully inside a large one.
975        let big = BoundingBox::new(0.0, 0.0, 100.0, 100.0);
976        let small = BoundingBox::new(10.0, 10.0, 20.0, 20.0);
977        assert!((ios_value(&big, &small) - 1.0).abs() < 1e-6); // fully contained
978        assert!(iou_value(&big, &small) < 0.05); // but IoU tiny
979    }
980
981    #[test]
982    fn jaccard_batch4_matches_scalar() {
983        let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
984        let boxes = [
985            BoundingBox::new(5.0, 5.0, 15.0, 15.0),   // overlap
986            BoundingBox::new(20.0, 20.0, 30.0, 30.0), // no overlap
987            BoundingBox::new(0.0, 0.0, 10.0, 10.0),   // identical
988            BoundingBox::new(8.0, 8.0, 18.0, 18.0),   // small overlap
989        ];
990        let iou_threshold = 0.1;
991        let batch = jaccard_batch4(&a, &boxes, iou_threshold);
992        for (i, b) in boxes.iter().enumerate() {
993            let scalar = jaccard(&a, b, iou_threshold);
994            assert_eq!(
995                batch[i], scalar,
996                "batch4 mismatch at {i}: batch={} scalar={}",
997                batch[i], scalar
998            );
999        }
1000    }
1001}