This commit is contained in:
Glenn Jocher 2019-04-26 23:25:00 +02:00
parent 96ea6a87cb
commit 84ebb5d143
1 changed files with 4 additions and 5 deletions

View File

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