From a5160b44caf9c67e18addd3e9c0c6651783ae067 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 23 Dec 2019 10:13:20 -0800 Subject: [PATCH] updates --- utils/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/utils/utils.py b/utils/utils.py index ef372ee7..b9180e13 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -528,8 +528,9 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru # Non-maximum suppression det_max = [] - for c in j.unique(): - dc = pred[j == c] # select class c + cls = pred[:, -1] + for c in cls.unique(): + dc = pred[cls == c] # select class c n = len(dc) if n == 1: det_max.append(dc) # No NMS required if only 1 prediction