updates
This commit is contained in:
parent
5263608d12
commit
96ea6a87cb
9
test.py
9
test.py
|
@ -124,16 +124,17 @@ def test(
|
|||
break
|
||||
|
||||
# Continue if predicted class not among image classes
|
||||
if pcls.item() not in tcls:
|
||||
m = (pcls == tcls_tensor).nonzero().view(-1) # matches
|
||||
if not any(m):
|
||||
continue
|
||||
|
||||
# Best iou, index between pred and targets
|
||||
iou, bi = bbox_iou(pbox, tbox).max(0)
|
||||
iou, bi = bbox_iou(pbox, tbox[m]).max(0)
|
||||
|
||||
# If iou > threshold and class is correct mark as correct
|
||||
if iou > iou_thres and bi not in detected: # and pcls == tcls[bi]:
|
||||
if iou > iou_thres and m[bi] not in detected: # and pcls == tcls[bi]:
|
||||
correct[i] = 1
|
||||
detected.append(bi)
|
||||
detected.append(m[bi])
|
||||
|
||||
# Append statistics (correct, conf, pcls, tcls)
|
||||
stats.append((correct, pred[:, 4].cpu(), pred[:, 6].cpu(), tcls))
|
||||
|
|
Loading…
Reference in New Issue