From 5545ff3545d886417b4eff12203d1af4d758cc10 Mon Sep 17 00:00:00 2001 From: Yonghye Kwon Date: Tue, 20 Dec 2022 01:19:14 +0900 Subject: [PATCH] Sort by confidence and remove excess boxes without explicit if (#10517) * sort by confidence and remove excess boxes without explicit if Signed-off-by: Yonghye Kwon * cleanup indexing boxes for remove excess boxes it is related to https://github.com/ultralytics/yolov5/pull/10502. Signed-off-by: Yonghye Kwon Co-authored-by: Glenn Jocher --- utils/general.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/utils/general.py b/utils/general.py index 70b6f64..0bbcb6e 100644 --- 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