From 597f239c577542a9fe272d18dc73ff12e1d08f18 Mon Sep 17 00:00:00 2001
From: BFlat <bombayapirat@gmail.com>
Date: Fri, 9 Oct 2026 16:16:19 +0700
Subject: [PATCH] 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 <noreply@anthropic.com>
---
 crates/segment/src/mediapipe.rs | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/crates/segment/src/mediapipe.rs b/crates/segment/src/mediapipe.rs
index eb4ab807..5798bc05 100644
--- a/crates/segment/src/mediapipe.rs
+++ b/crates/segment/src/mediapipe.rs
@@ -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).