diff --git a/test.py b/test.py index 163b7092..a334efcc 100644 --- a/test.py +++ b/test.py @@ -1,5 +1,6 @@ import argparse import json +import time from pathlib import Path from models import * @@ -47,6 +48,7 @@ def test( AP_accum, AP_accum_count = np.zeros(nC), np.zeros(nC) coco91class = coco80_to_coco91_class() for batch_i, (imgs, targets, paths, shapes) in enumerate(dataloader): + t = time.time() output = model(imgs.to(device)) output = non_max_suppression(output, conf_thres=conf_thres, nms_thres=nms_thres) @@ -128,13 +130,13 @@ def test( mean_P = np.mean(mP) # Print image mAP and running mean mAP - print(('%11s%11s' + '%11.3g' * 3) % (seen, dataloader.nF, mean_P, mean_R, mean_mAP)) + print(('%11s%11s' + '%11.3g' * 4 + 's') % + (seen, dataloader.nF, mean_P, mean_R, mean_mAP, time.time() - t)) # Print mAP per class print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:') - classes = load_classes(data_cfg_dict['names']) # Extracts class labels from file - for i, c in enumerate(classes): + for i, c in enumerate(load_classes(data_cfg_dict['names'])): print('%15s: %-.4f' % (c, AP_accum[i] / (AP_accum_count[i] + 1E-16))) # Save JSON diff --git a/utils/utils.py b/utils/utils.py index 7c4c0c81..e4982602 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -374,7 +374,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4): if prediction.is_cuda: unique_labels = unique_labels.cuda(prediction.device) - nms_style = 'MERGE' # 'OR' (default), 'AND', 'MERGE' (experimental) + nms_style = 'OR' # 'OR' (default), 'AND', 'MERGE' (experimental) for c in unique_labels: # Get the detections with class c dc = detections[detections[:, -1] == c]