updates

2019-12-23 10:11:12 -08:00 · 2019-12-23 10:11:12 -08:00 · a51d83df33
parent dd5ead5b1d
commit a51d83df33
1 changed files with 2 additions and 2 deletions
--- a/utils/utils.py
+++ b/utils/utils.py
@ -528,8 +528,8 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5, multi_cls=Tru

        # Non-maximum suppression
        det_max = []
-        for c in pred[:, -1].unique():
-            dc = pred[pred[:, -1] == c]  # select class c
+        for c in j.unique():
+            dc = pred[j == c]  # select class c
            n = len(dc)
            if n == 1:
                det_max.append(dc)  # No NMS required if only 1 prediction