From 84ebb5d143006cce60e13a141715e717634a3618 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 26 Apr 2019 23:25:00 +0200 Subject: [PATCH] updates --- test.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/test.py b/test.py index fe1a4128..c9516953 100644 --- a/test.py +++ b/test.py @@ -124,17 +124,16 @@ def test( break # Continue if predicted class not among image classes - m = (pcls == tcls_tensor).nonzero().view(-1) # matches - if not any(m): + if pcls.item() not in tcls: continue # 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 > 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 - detected.append(m[bi]) + detected.append(bi) # Append statistics (correct, conf, pcls, tcls) stats.append((correct, pred[:, 4].cpu(), pred[:, 6].cpu(), tcls))