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 use std::arch::aarch64::*;
518
519 let zero = vdupq_n_f32(0.0);
520 let iou_v = vdupq_n_f32(iou);
521
522 let a_xmin = vdupq_n_f32(a.xmin);
524 let a_ymin = vdupq_n_f32(a.ymin);
525 let a_xmax = vdupq_n_f32(a.xmax);
526 let a_ymax = vdupq_n_f32(a.ymax);
527 let area_a = vmulq_f32(vsubq_f32(a_xmax, a_xmin), vsubq_f32(a_ymax, a_ymin));
528
529 let b0 = vld1q_f32(&boxes[0].xmin as *const f32);
531 let b1 = vld1q_f32(&boxes[1].xmin as *const f32);
532 let b2 = vld1q_f32(&boxes[2].xmin as *const f32);
533 let b3 = vld1q_f32(&boxes[3].xmin as *const f32);
534
535 let t01_lo = vtrn1q_f32(b0, b1); let t01_hi = vtrn2q_f32(b0, b1); let t23_lo = vtrn1q_f32(b2, b3);
539 let t23_hi = vtrn2q_f32(b2, b3);
540
541 let b_xmin = vreinterpretq_f32_f64(vtrn1q_f64(
542 vreinterpretq_f64_f32(t01_lo),
543 vreinterpretq_f64_f32(t23_lo),
544 ));
545 let b_ymin = vreinterpretq_f32_f64(vtrn1q_f64(
546 vreinterpretq_f64_f32(t01_hi),
547 vreinterpretq_f64_f32(t23_hi),
548 ));
549 let b_xmax = vreinterpretq_f32_f64(vtrn2q_f64(
550 vreinterpretq_f64_f32(t01_lo),
551 vreinterpretq_f64_f32(t23_lo),
552 ));
553 let b_ymax = vreinterpretq_f32_f64(vtrn2q_f64(
554 vreinterpretq_f64_f32(t01_hi),
555 vreinterpretq_f64_f32(t23_hi),
556 ));
557
558 let left = vmaxq_f32(a_xmin, b_xmin);
560 let top = vmaxq_f32(a_ymin, b_ymin);
561 let right = vminq_f32(a_xmax, b_xmax);
562 let bottom = vminq_f32(a_ymax, b_ymax);
563 let w = vmaxq_f32(vsubq_f32(right, left), zero);
564 let h = vmaxq_f32(vsubq_f32(bottom, top), zero);
565 let intersection = vmulq_f32(w, h);
566
567 let area_b = vmulq_f32(vsubq_f32(b_xmax, b_xmin), vsubq_f32(b_ymax, b_ymin));
569
570 let union = vsubq_f32(vaddq_f32(area_a, area_b), intersection);
572
573 let iou_union = vmulq_f32(iou_v, union);
575 let mask = vcgtq_f32(intersection, iou_union);
576
577 [
579 vgetq_lane_u32(mask, 0) != 0,
580 vgetq_lane_u32(mask, 1) != 0,
581 vgetq_lane_u32(mask, 2) != 0,
582 vgetq_lane_u32(mask, 3) != 0,
583 ]
584}
585
586#[cfg(test)]
587mod tests {
588 use super::*;
589 use crate::BoundingBox;
590
591 fn make_nms_boxes_float(n: usize) -> Vec<DetectBox> {
593 (0..n)
594 .map(|i| DetectBox {
595 bbox: BoundingBox {
596 xmin: i as f32 * 100.0,
597 ymin: 0.0,
598 xmax: i as f32 * 100.0 + 10.0,
599 ymax: 10.0,
600 },
601 label: 0,
602 score: 1.0 - i as f32 * 0.01,
603 })
604 .collect()
605 }
606
607 #[test]
608 fn nms_float_max_det_matches_full_truncated() {
609 let boxes = make_nms_boxes_float(20);
610 let n = 5;
611 let full = nms_float(0.5, None, boxes.clone());
612 let capped = nms_float(0.5, Some(n), boxes);
613 assert_eq!(capped.len(), n);
614 for (f, c) in full[..n].iter().zip(capped.iter()) {
615 assert_eq!(f.bbox, c.bbox);
616 assert_eq!(f.score, c.score);
617 }
618 }
619
620 #[test]
621 fn nms_float_max_det_zero_returns_empty() {
622 let boxes = make_nms_boxes_float(10);
623 let result = nms_float(0.5, Some(0), boxes);
624 assert!(result.is_empty());
625 }
626
627 #[test]
628 fn nms_float_max_det_iou_ge_1_returns_sorted_truncated() {
629 let boxes = make_nms_boxes_float(10);
630 let result = nms_float(1.0, Some(3), boxes);
631 assert_eq!(result.len(), 3);
632 assert!(result[0].score >= result[1].score);
633 assert!(result[1].score >= result[2].score);
634 }
635
636 #[test]
637 fn nms_float_max_det_larger_than_input() {
638 let boxes = make_nms_boxes_float(5);
639 let full = nms_float(0.5, None, boxes.clone());
640 let capped = nms_float(0.5, Some(100), boxes);
641 assert_eq!(full.len(), capped.len());
642 }
643
644 #[test]
648 fn metric_values_match_modelpack_match() {
649 let a = BoundingBox::new(100.0, 100.0, 400.0, 300.0);
650 let b = BoundingBox::new(350.0, 100.0, 400.0, 300.0);
651 assert!((intersection_area(&a, &b) - 10000.0).abs() < 1e-3);
652 assert!((box_area(&a) - 60000.0).abs() < 1e-3);
653 assert!((box_area(&b) - 10000.0).abs() < 1e-3);
654 assert!((ios_value(&a, &b) - 1.0).abs() < 1e-6);
655 assert!((iou_value(&a, &b) - (1.0 / 6.0)).abs() < 1e-6);
656 }
657
658 #[test]
659 fn metric_values_disjoint_are_zero() {
660 let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
661 let b = BoundingBox::new(20.0, 20.0, 30.0, 30.0);
662 assert_eq!(intersection_area(&a, &b), 0.0);
663 assert_eq!(ios_value(&a, &b), 0.0);
664 assert_eq!(iou_value(&a, &b), 0.0);
665 }
666
667 #[test]
668 fn metric_values_zero_area_box_no_panic() {
669 let degenerate = BoundingBox::new(10.0, 10.0, 10.0, 20.0);
671 let other = BoundingBox::new(0.0, 0.0, 30.0, 30.0);
672 assert_eq!(box_area(°enerate), 0.0);
673 assert!(ios_value(°enerate, &other).is_finite());
674 assert!(iou_value(°enerate, &other).is_finite());
675 }
676
677 #[test]
678 fn ios_high_where_iou_low_for_contained_box() {
679 let big = BoundingBox::new(0.0, 0.0, 100.0, 100.0);
681 let small = BoundingBox::new(10.0, 10.0, 20.0, 20.0);
682 assert!((ios_value(&big, &small) - 1.0).abs() < 1e-6); assert!(iou_value(&big, &small) < 0.05); }
685
686 #[test]
687 fn jaccard_batch4_matches_scalar() {
688 let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
689 let boxes = [
690 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), ];
695 let iou_threshold = 0.1;
696 let batch = jaccard_batch4(&a, &boxes, iou_threshold);
697 for (i, b) in boxes.iter().enumerate() {
698 let scalar = jaccard(&a, b, iou_threshold);
699 assert_eq!(
700 batch[i], scalar,
701 "batch4 mismatch at {i}: batch={} scalar={}",
702 batch[i], scalar
703 );
704 }
705 }
706}