updates

utils/datasets.py

@@ -128,8 +128,7 @@ class LoadImagesAndLabels:  # for training
             # Fixed-Scale YOLO Training
             height = self.height
 
-        img_all = []
-        labels_all = []
+        img_all, labels_all, img_paths, img_shapes = [], [], [], []
         for index, files_index in enumerate(range(ia, ib)):
             img_path = self.img_files[self.shuffled_vector[files_index]]
             label_path = self.label_files[self.shuffled_vector[files_index]]
@@ -210,13 +209,15 @@ class LoadImagesAndLabels:  # for training
 
             img_all.append(img)
             labels_all.append(torch.from_numpy(labels))
+            img_paths.append(img_path)
+            img_shapes.append((h, w))
 
         # Normalize
         img_all = np.stack(img_all)[:, :, :, ::-1].transpose(0, 3, 1, 2)  # BGR to RGB and cv2 to pytorch
         img_all = np.ascontiguousarray(img_all, dtype=np.float32)
         img_all /= 255.0
 
-        return torch.from_numpy(img_all), labels_all
+        return torch.from_numpy(img_all), labels_all, img_paths, img_shapes
 
     def __len__(self):
         return self.nB  # number of batches

utils/torch_utils.py

@@ -20,5 +20,5 @@ def select_device(force_cpu=False):
             torch.cuda.set_device(0)  # OPTIONAL: Set your GPU if multiple available
             # print('Using ', torch.cuda.device_count(), ' GPUs')
 
-    print('Using %s  %s\n' % (device.type, torch.cuda.get_device_properties(0) if cuda else ''))
+    print('Using %s %s\n' % (device.type, torch.cuda.get_device_properties(0) if cuda else ''))
     return device

utils/utils.py

@@ -49,6 +49,14 @@ 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
+    # 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]
+
+
 def plot_one_box(x, img, color=None, label=None, line_thickness=None):  # Plots one bounding box on image img
     tl = line_thickness or round(0.002 * max(img.shape[0:2])) + 1  # line thickness
     color = color or [random.randint(0, 255) for _ in range(3)]
@@ -99,7 +107,7 @@ def scale_coords(img_size, coords, img0_shape):
     coords[:, [0, 2]] -= pad_x
     coords[:, [1, 3]] -= pad_y
     coords[:, :4] /= gain
-    coords[:, :4] = torch.round(torch.clamp(coords[:, :4], min=0))
+    coords[:, :4] = torch.clamp(coords[:, :4], min=0)
     return coords