Sort by confidence and remove excess boxes without explicit if (ultra…

…lytics#10517) * sort by confidence and remove excess boxes without explicit if Signed-off-by: Yonghye Kwon <[email protected]> * cleanup indexing boxes for remove excess boxes it is related to ultralytics#10502. Signed-off-by: Yonghye Kwon <[email protected]> Co-authored-by: Glenn Jocher <[email protected]>
ktho22 · Dec 19, 2022 · 5545ff3 · 5545ff3
1 parent 10e93d2
commit 5545ff3
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/utils/general.py b/utils/general.py
@@ -969,10 +969,7 @@ def non_max_suppression(
         n = x.shape[0]  # number of boxes
         if not n:  # no boxes
             continue
-        elif n > max_nms:  # excess boxes
-            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence
-        else:
-            x = x[x[:, 4].argsort(descending=True)]  # sort by confidence
+        x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence and remove excess boxes
 
         # Batched NMS
         c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes