Update test.py with set() (#806)

2020-08-21 15:03:03 -07:00 · 2020-08-21 15:03:03 -07:00 · c171e45cda
parent b42e8a531b
commit c171e45cda
1 changed files with 3 additions and 1 deletions
--- a/test.py
+++ b/test.py
@ -172,9 +172,11 @@ def test(data,
                        ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1)  # best ious, indices

                        # Append detections
+                        detected_set = set()
                        for j in (ious > iouv[0]).nonzero(as_tuple=False):
                            d = ti[i[j]]  # detected target
-                            if d not in detected:
+                            if d.item() not in detected_set:
+                                detected_set.add(d.item())
                                detected.append(d)
                                correct[pi[j]] = ious[j] > iouv  # iou_thres is 1xn
                                if len(detected) == nl:  # all targets already located in image