diff --git a/test.py b/test.py index a848369d..f8739866 100644 --- a/test.py +++ b/test.py @@ -37,18 +37,15 @@ def test( model.to(device).eval() # Get dataloader - # dataloader = torch.utils.data.DataLoader(LoadImagesAndLabels(test_path), batch_size=batch_size) # pytorch + # dataloader = torch.utils.data.DataLoader(LoadImagesAndLabels(test_path), batch_size=batch_size) dataloader = LoadImagesAndLabels(test_path, batch_size=batch_size, img_size=img_size) - # Create JSON - jdict = [] - float3 = lambda x: float(format(x, '.3f')) # print json to 3 decimals - # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... - mean_mAP, mean_R, mean_P, seen = 0.0, 0.0, 0.0, 0 print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP')) - outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class = [], [], [], [], [], [], [], [] + outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class, jdict = \ + [], [], [], [], [], [], [], [], [] 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): output = model(imgs.to(device)) output = non_max_suppression(output, conf_thres=conf_thres, nms_thres=nms_thres) @@ -67,18 +64,18 @@ def test( detections = detections.cpu().numpy() detections = detections[np.argsort(-detections[:, 4])] - # Save JSON if save_json: - # rescale box to original image size, top left origin - box = torch.from_numpy(detections[:, :4]).clone() # x1y1x2y2 - scale_coords(img_size, box, shapes[si]) - box = xyxy2xywh(box) - box[:, :2] -= box[:, 2:] / 2 # origin center to corner + # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... + box = torch.from_numpy(detections[:, :4]).clone() # xyxy + scale_coords(img_size, box, shapes[si]) # to original shape + box = xyxy2xywh(box) # xywh + box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner + # add to json dictionary for di, d in enumerate(detections): - jdict.append({ # add to json dictionary + jdict.append({ 'image_id': int(Path(paths[si]).stem.split('_')[-1]), - 'category_id': darknet2coco_class(int(d[6])), + 'category_id': coco91class(int(d[6])), 'bbox': [float3(x) for x in box[di]], 'score': float3(d[4] * d[5]) }) diff --git a/utils/utils.py b/utils/utils.py index fa9df88a..ae52c4da 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -12,6 +12,10 @@ torch.set_printoptions(linewidth=1320, precision=5, profile='long') np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 +def float3(x): # format floats to 3 decimals + return float(format(x, '.3f')) + + def init_seeds(seed=0): random.seed(seed) np.random.seed(seed) @@ -49,12 +53,12 @@ def coco_class_weights(): # frequency of each class in coco train2014 return weights -def darknet2coco_class(c): # returns the coco class for each darknet class +def coco80_to_coco91_class(): # returns the coco class for each darknet class # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/ a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') x = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco - return x[c] + return x def plot_one_box(x, img, color=None, label=None, line_thickness=None): # Plots one bounding box on image img