1use crate::{arg_max, BBoxTypeTrait, BoundingBox, DetectBox};
21use ndarray::{
22 parallel::prelude::{IntoParallelIterator, ParallelIterator as _},
23 Array1, ArrayView2, Zip,
24};
25use num_traits::{AsPrimitive, Float};
26use rayon::slice::ParallelSliceMut;
27
28pub fn postprocess_boxes_float<
32 B: BBoxTypeTrait,
33 BOX: Float + AsPrimitive<f32> + Send + Sync,
34 SCORE: Float + AsPrimitive<f32> + Send + Sync,
35>(
36 threshold: SCORE,
37 boxes: ArrayView2<BOX>,
38 scores: ArrayView2<SCORE>,
39) -> Vec<DetectBox> {
40 assert_eq!(scores.dim().0, boxes.dim().0);
41 assert_eq!(boxes.dim().1, 4);
42 Zip::from(scores.rows())
43 .and(boxes.rows())
44 .into_par_iter()
45 .filter_map(|(score, bbox)| {
46 let (score_, label) = arg_max(score);
47 if score_ < threshold {
48 return None;
49 }
50
51 let bbox = B::ndarray_to_xyxy_float(bbox);
52 Some(DetectBox {
53 label,
54 score: score_.as_(),
55 bbox: bbox.into(),
56 })
57 })
58 .collect()
59}
60
61pub fn postprocess_boxes_index_float<
68 B: BBoxTypeTrait,
69 BOX: Float + AsPrimitive<f32> + Send + Sync,
70 SCORE: Float + AsPrimitive<f32> + Send + Sync,
71>(
72 threshold: SCORE,
73 boxes: ArrayView2<BOX>,
74 scores: ArrayView2<SCORE>,
75) -> Vec<(DetectBox, usize)> {
76 assert_eq!(scores.dim().0, boxes.dim().0);
77 assert_eq!(boxes.dim().1, 4);
78 let indices: Array1<usize> = (0..boxes.dim().0).collect();
79 Zip::from(scores.rows())
80 .and(boxes.rows())
81 .and(&indices)
82 .into_par_iter()
83 .filter_map(|(score, bbox, i)| {
84 let (score_, label) = arg_max(score);
85 if score_ < threshold {
86 return None;
87 }
88
89 let bbox = B::ndarray_to_xyxy_float(bbox);
90 Some((
91 DetectBox {
92 label,
93 score: score_.as_(),
94 bbox: bbox.into(),
95 },
96 *i,
97 ))
98 })
99 .collect()
100}
101
102pub fn postprocess_boxes_multilabel_index_float<
118 B: BBoxTypeTrait,
119 BOX: Float + AsPrimitive<f32> + Send + Sync,
120 SCORE: Float + AsPrimitive<f32> + Send + Sync,
121>(
122 threshold: SCORE,
123 boxes: ArrayView2<BOX>,
124 scores: ArrayView2<SCORE>,
125) -> Vec<(DetectBox, usize)> {
126 assert_eq!(scores.dim().0, boxes.dim().0);
127 assert_eq!(boxes.dim().1, 4);
128 let n = boxes.dim().0;
129 let indices: Array1<usize> = (0..n).collect();
130 Zip::from(scores.rows())
131 .and(boxes.rows())
132 .and(&indices)
133 .into_par_iter()
134 .flat_map(|(score_row, bbox_row, &anchor_idx)| {
135 let bbox = B::ndarray_to_xyxy_float(bbox_row);
137 let bbox: crate::BoundingBox = bbox.into();
138 score_row
139 .iter()
140 .enumerate()
141 .filter(|(_, &s)| s >= threshold)
142 .map(move |(c, &s)| {
143 (
144 DetectBox {
145 label: c,
146 score: s.as_(),
147 bbox,
148 },
149 anchor_idx,
150 )
151 })
152 .collect::<Vec<_>>()
153 })
154 .collect()
155}
156
157#[must_use]
165pub fn nms_float(iou: f32, max_det: Option<usize>, mut boxes: Vec<DetectBox>) -> Vec<DetectBox> {
166 boxes.par_sort_by(|a, b| b.score.total_cmp(&a.score));
169
170 if iou >= 1.0 {
173 return match max_det {
174 Some(n) => {
175 boxes.truncate(n);
176 boxes
177 }
178 None => boxes,
179 };
180 }
181
182 let cap = max_det.unwrap_or(usize::MAX);
183 let mut survivors: usize = 0;
184
185 for i in 0..boxes.len() {
187 if boxes[i].score < 0.0 {
188 continue;
190 }
191 for j in (i + 1)..boxes.len() {
192 if boxes[j].score < 0.0 {
195 continue;
197 }
198 if jaccard(&boxes[j].bbox, &boxes[i].bbox, iou) {
199 boxes[j].score = -1.0;
201 }
202 }
203
204 survivors += 1;
209 if survivors >= cap {
210 break;
211 }
212 }
213 boxes
216 .into_iter()
217 .filter(|b| b.score >= 0.0)
218 .take(cap)
219 .collect()
220}
221
222#[must_use]
228pub fn nms_extra_float<E: Send + Sync>(
229 iou: f32,
230 max_det: Option<usize>,
231 mut boxes: Vec<(DetectBox, E)>,
232) -> Vec<(DetectBox, E)> {
233 boxes.par_sort_by(|a, b| b.0.score.total_cmp(&a.0.score));
236
237 if iou >= 1.0 {
240 return match max_det {
241 Some(n) => {
242 boxes.truncate(n);
243 boxes
244 }
245 None => boxes,
246 };
247 }
248
249 let cap = max_det.unwrap_or(usize::MAX);
250 let mut survivors: usize = 0;
251
252 for i in 0..boxes.len() {
254 if boxes[i].0.score < 0.0 {
255 continue;
257 }
258 for j in (i + 1)..boxes.len() {
259 if boxes[j].0.score < 0.0 {
262 continue;
264 }
265 if jaccard(&boxes[j].0.bbox, &boxes[i].0.bbox, iou) {
266 boxes[j].0.score = -1.0;
268 }
269 }
270 survivors += 1;
271 if survivors >= cap {
272 break;
273 }
274 }
275
276 boxes
279 .into_iter()
280 .filter(|b| b.0.score >= 0.0)
281 .take(cap)
282 .collect()
283}
284
285#[must_use]
311pub fn nms_class_aware_float(
312 iou: f32,
313 max_det: Option<usize>,
314 mut boxes: Vec<DetectBox>,
315) -> Vec<DetectBox> {
316 boxes.par_sort_by(|a, b| b.score.total_cmp(&a.score));
317
318 if iou >= 1.0 {
319 return match max_det {
320 Some(n) => {
321 boxes.truncate(n);
322 boxes
323 }
324 None => boxes,
325 };
326 }
327
328 let cap = max_det.unwrap_or(usize::MAX);
329 let mut survivors: usize = 0;
330
331 for i in 0..boxes.len() {
332 if boxes[i].score < 0.0 {
333 continue;
334 }
335 for j in (i + 1)..boxes.len() {
336 if boxes[j].score < 0.0 {
337 continue;
338 }
339 if boxes[j].label == boxes[i].label && jaccard(&boxes[j].bbox, &boxes[i].bbox, iou) {
341 boxes[j].score = -1.0;
342 }
343 }
344 survivors += 1;
345 if survivors >= cap {
346 break;
347 }
348 }
349 boxes
350 .into_iter()
351 .filter(|b| b.score >= 0.0)
352 .take(cap)
353 .collect()
354}
355
356#[must_use]
361pub fn nms_extra_class_aware_float<E: Send + Sync>(
362 iou: f32,
363 max_det: Option<usize>,
364 mut boxes: Vec<(DetectBox, E)>,
365) -> Vec<(DetectBox, E)> {
366 boxes.par_sort_by(|a, b| b.0.score.total_cmp(&a.0.score));
367
368 if iou >= 1.0 {
371 return match max_det {
372 Some(n) => {
373 boxes.truncate(n);
374 boxes
375 }
376 None => boxes,
377 };
378 }
379
380 let cap = max_det.unwrap_or(usize::MAX);
381 let mut survivors: usize = 0;
382
383 for i in 0..boxes.len() {
384 if boxes[i].0.score < 0.0 {
385 continue;
386 }
387 for j in (i + 1)..boxes.len() {
388 if boxes[j].0.score < 0.0 {
389 continue;
390 }
391 if boxes[j].0.label == boxes[i].0.label
393 && jaccard(&boxes[j].0.bbox, &boxes[i].0.bbox, iou)
394 {
395 boxes[j].0.score = -1.0;
396 }
397 }
398 survivors += 1;
399 if survivors >= cap {
400 break;
401 }
402 }
403 boxes
404 .into_iter()
405 .filter(|b| b.0.score >= 0.0)
406 .take(cap)
407 .collect()
408}
409
410#[must_use]
413#[inline]
414pub fn intersection_area(a: &BoundingBox, b: &BoundingBox) -> f32 {
415 let left = a.xmin.max(b.xmin);
416 let top = a.ymin.max(b.ymin);
417 let right = a.xmax.min(b.xmax);
418 let bottom = a.ymax.min(b.ymax);
419 (right - left).max(0.0) * (bottom - top).max(0.0)
420}
421
422#[must_use]
424#[inline]
425pub fn box_area(b: &BoundingBox) -> f32 {
426 (b.xmax - b.xmin).max(0.0) * (b.ymax - b.ymin).max(0.0)
427}
428
429#[must_use]
432#[inline]
433pub fn iou_value(a: &BoundingBox, b: &BoundingBox) -> f32 {
434 let inter = intersection_area(a, b);
435 let union = (box_area(a) + box_area(b) - inter).max(1e-9);
436 inter / union
437}
438
439#[must_use]
443#[inline]
444pub fn ios_value(a: &BoundingBox, b: &BoundingBox) -> f32 {
445 let inter = intersection_area(a, b);
446 let denom = box_area(a).min(box_area(b)).max(1e-9);
447 inter / denom
448}
449
450pub fn jaccard(a: &BoundingBox, b: &BoundingBox, iou: f32) -> bool {
468 let left = a.xmin.max(b.xmin);
469 let top = a.ymin.max(b.ymin);
470 let right = a.xmax.min(b.xmax);
471 let bottom = a.ymax.min(b.ymax);
472
473 let intersection = (right - left).max(0.0) * (bottom - top).max(0.0);
474 let area_a = (a.xmax - a.xmin) * (a.ymax - a.ymin);
475 let area_b = (b.xmax - b.xmin) * (b.ymax - b.ymin);
476
477 let union = area_a + area_b - intersection;
479
480 intersection > iou * union
481}
482
483#[inline]
492pub fn jaccard_batch4(a: &BoundingBox, boxes: &[BoundingBox; 4], iou: f32) -> [bool; 4] {
493 #[cfg(target_arch = "aarch64")]
494 {
495 unsafe { jaccard_batch4_neon(a, boxes, iou) }
497 }
498 #[cfg(not(target_arch = "aarch64"))]
499 {
500 [
501 jaccard(a, &boxes[0], iou),
502 jaccard(a, &boxes[1], iou),
503 jaccard(a, &boxes[2], iou),
504 jaccard(a, &boxes[3], iou),
505 ]
506 }
507}
508
509#[cfg(target_arch = "aarch64")]
515#[target_feature(enable = "neon")]
516unsafe fn jaccard_batch4_neon(a: &BoundingBox, boxes: &[BoundingBox; 4], iou: f32) -> [bool; 4] {
517 unsafe {
518 use std::arch::aarch64::*; let zero = vdupq_n_f32(0.0);
521 let iou_v = vdupq_n_f32(iou);
522
523 let a_xmin = vdupq_n_f32(a.xmin);
525 let a_ymin = vdupq_n_f32(a.ymin);
526 let a_xmax = vdupq_n_f32(a.xmax);
527 let a_ymax = vdupq_n_f32(a.ymax);
528 let area_a = vmulq_f32(vsubq_f32(a_xmax, a_xmin), vsubq_f32(a_ymax, a_ymin));
529
530 let b0 = vld1q_f32(&boxes[0].xmin as *const f32);
532 let b1 = vld1q_f32(&boxes[1].xmin as *const f32);
533 let b2 = vld1q_f32(&boxes[2].xmin as *const f32);
534 let b3 = vld1q_f32(&boxes[3].xmin as *const f32);
535
536 let t01_lo = vtrn1q_f32(b0, b1); let t01_hi = vtrn2q_f32(b0, b1); let t23_lo = vtrn1q_f32(b2, b3);
540 let t23_hi = vtrn2q_f32(b2, b3);
541
542 let b_xmin = vreinterpretq_f32_f64(vtrn1q_f64(
543 vreinterpretq_f64_f32(t01_lo),
544 vreinterpretq_f64_f32(t23_lo),
545 ));
546 let b_ymin = vreinterpretq_f32_f64(vtrn1q_f64(
547 vreinterpretq_f64_f32(t01_hi),
548 vreinterpretq_f64_f32(t23_hi),
549 ));
550 let b_xmax = vreinterpretq_f32_f64(vtrn2q_f64(
551 vreinterpretq_f64_f32(t01_lo),
552 vreinterpretq_f64_f32(t23_lo),
553 ));
554 let b_ymax = vreinterpretq_f32_f64(vtrn2q_f64(
555 vreinterpretq_f64_f32(t01_hi),
556 vreinterpretq_f64_f32(t23_hi),
557 ));
558
559 let left = vmaxq_f32(a_xmin, b_xmin);
561 let top = vmaxq_f32(a_ymin, b_ymin);
562 let right = vminq_f32(a_xmax, b_xmax);
563 let bottom = vminq_f32(a_ymax, b_ymax);
564 let w = vmaxq_f32(vsubq_f32(right, left), zero);
565 let h = vmaxq_f32(vsubq_f32(bottom, top), zero);
566 let intersection = vmulq_f32(w, h);
567
568 let area_b = vmulq_f32(vsubq_f32(b_xmax, b_xmin), vsubq_f32(b_ymax, b_ymin));
570
571 let union = vsubq_f32(vaddq_f32(area_a, area_b), intersection);
573
574 let iou_union = vmulq_f32(iou_v, union);
576 let mask = vcgtq_f32(intersection, iou_union);
577
578 [
580 vgetq_lane_u32(mask, 0) != 0,
581 vgetq_lane_u32(mask, 1) != 0,
582 vgetq_lane_u32(mask, 2) != 0,
583 vgetq_lane_u32(mask, 3) != 0,
584 ]
585 }
586}
587
588#[cfg(test)]
589mod tests {
590 use super::*;
591 use crate::BoundingBox;
592
593 fn make_nms_boxes_float(n: usize) -> Vec<DetectBox> {
595 (0..n)
596 .map(|i| DetectBox {
597 bbox: BoundingBox {
598 xmin: i as f32 * 100.0,
599 ymin: 0.0,
600 xmax: i as f32 * 100.0 + 10.0,
601 ymax: 10.0,
602 },
603 label: 0,
604 score: 1.0 - i as f32 * 0.01,
605 })
606 .collect()
607 }
608
609 #[test]
610 fn nms_float_max_det_matches_full_truncated() {
611 let boxes = make_nms_boxes_float(20);
612 let n = 5;
613 let full = nms_float(0.5, None, boxes.clone());
614 let capped = nms_float(0.5, Some(n), boxes);
615 assert_eq!(capped.len(), n);
616 for (f, c) in full[..n].iter().zip(capped.iter()) {
617 assert_eq!(f.bbox, c.bbox);
618 assert_eq!(f.score, c.score);
619 }
620 }
621
622 #[test]
623 fn nms_float_max_det_zero_returns_empty() {
624 let boxes = make_nms_boxes_float(10);
625 let result = nms_float(0.5, Some(0), boxes);
626 assert!(result.is_empty());
627 }
628
629 #[test]
630 fn nms_float_max_det_iou_ge_1_returns_sorted_truncated() {
631 let boxes = make_nms_boxes_float(10);
632 let result = nms_float(1.0, Some(3), boxes);
633 assert_eq!(result.len(), 3);
634 assert!(result[0].score >= result[1].score);
635 assert!(result[1].score >= result[2].score);
636 }
637
638 #[test]
639 fn nms_float_max_det_larger_than_input() {
640 let boxes = make_nms_boxes_float(5);
641 let full = nms_float(0.5, None, boxes.clone());
642 let capped = nms_float(0.5, Some(100), boxes);
643 assert_eq!(full.len(), capped.len());
644 }
645
646 #[test]
650 fn metric_values_match_modelpack_match() {
651 let a = BoundingBox::new(100.0, 100.0, 400.0, 300.0);
652 let b = BoundingBox::new(350.0, 100.0, 400.0, 300.0);
653 assert!((intersection_area(&a, &b) - 10000.0).abs() < 1e-3);
654 assert!((box_area(&a) - 60000.0).abs() < 1e-3);
655 assert!((box_area(&b) - 10000.0).abs() < 1e-3);
656 assert!((ios_value(&a, &b) - 1.0).abs() < 1e-6);
657 assert!((iou_value(&a, &b) - (1.0 / 6.0)).abs() < 1e-6);
658 }
659
660 #[test]
661 fn metric_values_disjoint_are_zero() {
662 let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
663 let b = BoundingBox::new(20.0, 20.0, 30.0, 30.0);
664 assert_eq!(intersection_area(&a, &b), 0.0);
665 assert_eq!(ios_value(&a, &b), 0.0);
666 assert_eq!(iou_value(&a, &b), 0.0);
667 }
668
669 #[test]
670 fn metric_values_zero_area_box_no_panic() {
671 let degenerate = BoundingBox::new(10.0, 10.0, 10.0, 20.0);
673 let other = BoundingBox::new(0.0, 0.0, 30.0, 30.0);
674 assert_eq!(box_area(°enerate), 0.0);
675 assert!(ios_value(°enerate, &other).is_finite());
676 assert!(iou_value(°enerate, &other).is_finite());
677 }
678
679 #[test]
680 fn ios_high_where_iou_low_for_contained_box() {
681 let big = BoundingBox::new(0.0, 0.0, 100.0, 100.0);
683 let small = BoundingBox::new(10.0, 10.0, 20.0, 20.0);
684 assert!((ios_value(&big, &small) - 1.0).abs() < 1e-6); assert!(iou_value(&big, &small) < 0.05); }
687
688 #[test]
689 fn jaccard_batch4_matches_scalar() {
690 let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
691 let boxes = [
692 BoundingBox::new(5.0, 5.0, 15.0, 15.0), BoundingBox::new(20.0, 20.0, 30.0, 30.0), BoundingBox::new(0.0, 0.0, 10.0, 10.0), BoundingBox::new(8.0, 8.0, 18.0, 18.0), ];
697 let iou_threshold = 0.1;
698 let batch = jaccard_batch4(&a, &boxes, iou_threshold);
699 for (i, b) in boxes.iter().enumerate() {
700 let scalar = jaccard(&a, b, iou_threshold);
701 assert_eq!(
702 batch[i], scalar,
703 "batch4 mismatch at {i}: batch={} scalar={}",
704 batch[i], scalar
705 );
706 }
707 }
708}