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

This commit is contained in:
Glenn Jocher committed 2020-08-21 15:03:03 -07:00
1 parent b42e8a531b
commit c171e45cda
1 file changed
+3 -1
+3 -1
View File
@@ -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