updates

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])
                     })

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