Fix #367: weighted_nms terminates on NaN or zero-area boxes

A degenerate detection box (any NaN coordinate, or zero area) has an
overlap of 0 with itself, so `top` was partitioned into the remainder and
the loop never shrank `dets`. `top` now always joins its own cluster, so
every pass removes at least one detection. Well-formed boxes behave as
before (their self-overlap is 1).

Reported with the fix by KazW.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
BFlatandClaude Opus 5.5 committed 2026-10-09 16:16:19 +07:00
1 parent cae67546dd
commit 597f239c57
1 file changed
+16 -1
+16 -1
View File
@@ -107,8 +107,11 @@ fn weighted_nms(mut dets: Vec<Detection>) -> Vec<Detection> {
dets.sort_by(|a, b| b.score.total_cmp(&a.score));
let mut out = vec![];
while let Some(top) = dets.first().copied() {
let (cluster, rest): (Vec<Detection>, Vec<Detection>) = dets.into_iter().partition(|d| iou(&d.bbox, &top.bbox) > NMS_IOU);
// `top` always joins its own cluster: its overlap with itself is 0 for a degenerate box
// (NaN or zero area), so partitioning it with the rest would never shrink `dets`.
let (mut cluster, rest): (Vec<Detection>, Vec<Detection>) = dets.into_iter().skip(1).partition(|d| iou(&d.bbox, &top.bbox) > NMS_IOU);
dets = rest;
cluster.insert(0, top);
let total: f32 = cluster.iter().map(|d| d.score).sum();
if total <= 0.0 {
out.push(top);
@@ -283,6 +286,18 @@ mod tests {
assert_eq!(out[1].bbox[0], 50.0);
}
/// A degenerate box (NaN or zero area) overlaps nothing, not even itself: each still becomes
/// its own output instead of looping forever (#367).
#[test]
fn nms_terminates_on_degenerate_boxes() {
let d = Detection { bbox: [f32::NAN; 4], keypoints: [[f32::NAN; 2]; 6], score: 0.7 };
assert_eq!(weighted_nms(vec![d, d, d]).len(), 3);
let z = Detection { bbox: [0.0; 4], keypoints: [[0.0; 2]; 6], score: 0.0 };
assert_eq!(weighted_nms(vec![z, z]).len(), 2);
let n = Detection { bbox: [0.0; 4], keypoints: [[0.0; 2]; 6], score: f32::NAN };
assert_eq!(weighted_nms(vec![n, z, d]).len(), 3);
}
/// The official bundle (`EFFECTCRAFT_FACE_LANDMARKER` = path to `face_landmarker.task`) on a
/// fixed input matches TensorFlow Lite's own interpreter (`ai-edge-litert` 2.2.0, run once
/// to record these sums).