Operation Coco is underway

Author: amdegroot

data/__init__.py

@@ -5,7 +5,6 @@
 import cv2
 import numpy as np
 
-
 def detection_collate(batch):
     """Custom collate fn for dealing with batches of images that have a different
     number of associated object annotations (bounding boxes).

data/coco.py

@@ -1,14 +1,14 @@
 from .config import HOME
 import os
-import os.path
+import os.path as osp
 import sys
 import torch
 import torch.utils.data as data
 import torchvision.transforms as transforms
 import cv2
 import numpy as np
 
-COCO_ROOT = os.path.join(HOME, 'data/coco/')
+COCO_ROOT = osp.join(HOME, 'data/coco/')
 IMAGES = 'images'
 ANNOTATIONS = 'annotations'
 COCO_API = 'PythonAPI'
@@ -34,6 +34,8 @@ class COCOAnnotationTransform(object):
     """Transforms a COCO annotation into a Tensor of bbox coords and label index
     Initilized with a dictionary lookup of classnames to indexes
     """
+    def __init__(self):
+        self.label_map = get_label_map(osp.join(COCO_ROOT, 'coco_labels.txt'))
 
     def __call__(self, target, width, height):
         """
@@ -51,10 +53,13 @@ def __call__(self, target, width, height):
                 bbox = obj['bbox']
                 bbox[2] += bbox[0]
                 bbox[3] += bbox[1]
-                label_idx = obj['category_id']
+                label_idx = self.label_map[obj['category_id']] - 1
                 final_box = list(np.array(bbox)/scale)
                 final_box.append(label_idx)
                 res += [final_box]  # [xmin, ymin, xmax, ymax, label_idx]
+            else:
+                print("no bbox problem!")
+
         return res  # [[xmin, ymin, xmax, ymax, label_idx], ... ]
 
 
@@ -70,16 +75,16 @@ class COCODetection(data.Dataset):
     """
 
     def __init__(self, root, image_set, transform=None,
-                 target_transform=None, dataset_name='COCO2014'):
-        sys.path.append(os.path.join(root, COCO_API))
+                 target_transform=None):
+        sys.path.append(osp.join(root, COCO_API))
         from pycocotools.coco import COCO
-        self.root = os.path.join(root, IMAGES, image_set)
-        self.coco = COCO(os.path.join(root, ANNOTATIONS,
-                                      INSTANCES_SET.format(image_set)))
-        self.ids = list(self.coco.imgs.keys())
+        self.root = osp.join(root, IMAGES, image_set)
+        self.coco = COCO(osp.join(root, ANNOTATIONS,
+                                  INSTANCES_SET.format(image_set)))
+        self.ids = list(self.coco.imgToAnns.keys())
         self.transform = transform
         self.target_transform = target_transform
-        self.name = dataset_name
+        self.name = 'MS COCO ' + image_set
 
     def __getitem__(self, index):
         """
@@ -104,11 +109,14 @@ def pull_item(self, index):
                    target is the object returned by ``coco.loadAnns``.
         """
         img_id = self.ids[index]
+        target = self.coco.imgToAnns[img_id]
         ann_ids = self.coco.getAnnIds(imgIds=img_id)
+
         target = self.coco.loadAnns(ann_ids)
-        path = self.coco.loadImgs(img_id)[0]['file_name']
-        img = cv2.imread(os.path.join(self.root, path))
-        height, width, channels = img.shape
+        path = osp.join(self.root, self.coco.loadImgs(img_id)[0]['file_name'])
+        assert osp.exists(path), 'Image path does not exist: {}'.format(path)
+        img = cv2.imread(osp.join(self.root, path))
+        height, width, _ = img.shape
         if self.target_transform is not None:
             target = self.target_transform(target, width, height)
         if self.transform is not None:
@@ -117,7 +125,7 @@ def pull_item(self, index):
                                                 target[:, 4])
             # to rgb
             img = img[:, :, (2, 1, 0)]
-            # img = img.transpose(2, 0, 1)
+
             target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
         return torch.from_numpy(img).permute(2, 0, 1), target, height, width
 
@@ -134,7 +142,7 @@ def pull_image(self, index):
         '''
         img_id = self.ids[index]
         path = self.coco.loadImgs(img_id)[0]['file_name']
-        return cv2.imread(os.path.join(self.root, path), cv2.IMREAD_COLOR)
+        return cv2.imread(osp.join(self.root, path), cv2.IMREAD_COLOR)
 
     def pull_anno(self, index):
         '''Returns the original annotation of image at index
@@ -161,3 +169,12 @@ def __repr__(self):
         tmp = '    Target Transforms (if any): '
         fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
         return fmt_str
+
+
+def get_label_map(label_file):
+    label_map = {}
+    labels = open(label_file, 'r')
+    for line in labels:
+        ids = line.split(',')
+        label_map[int(ids[0])] = int(ids[1])
+    return label_map

layers/box_utils.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import torch
 
+
 def point_form(boxes):
     """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
     representation for comparison to point form ground truth data.

layers/modules/multibox_loss.py

@@ -54,7 +54,7 @@ def forward(self, predictions, targets):
                 loc shape: torch.size(batch_size,num_priors,4)
                 priors shape: torch.size(num_priors,4)
 
-            ground_truth (tensor): Ground truth boxes and labels for a batch,
+            targets (tensor): Ground truth boxes and labels for a batch,
                 shape: [batch_size,num_objs,5] (last idx is the label).
         """
         loc_data, conf_data, priors = predictions
@@ -91,12 +91,11 @@ def forward(self, predictions, targets):
 
         # Compute max conf across batch for hard negative mining
         batch_conf = conf_data.view(-1, self.num_classes)
-
         loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))
 
         # Hard Negative Mining
-        loss_c = loss_c.view(num, -1)
         loss_c[pos] = 0  # filter out pos boxes for now
+        loss_c = loss_c.view(num, -1)
         _, loss_idx = loss_c.sort(1, descending=True)
         _, idx_rank = loss_idx.sort(1)
         num_pos = pos.long().sum(1, keepdim=True)

train.py

@@ -34,8 +34,8 @@ def str2bool(v):
 parser.add_argument('--weight_decay', default=5e-4, type=float, help='Weight decay for SGD')
 parser.add_argument('--gamma', default=0.1, type=float, help='Gamma update for SGD')
 parser.add_argument('--log_iters', default=True, type=bool, help='Print the loss at each iteration')
-parser.add_argument('--visdom', default=True, type=str2bool, help='Use visdom for loss visualization')
-parser.add_argument('--send_images_to_visdom', type=str2bool, default=True, help='Sample a random image from every 10th batch, send it to visdom after augmentations step')
+parser.add_argument('--visdom', default=False, type=str2bool, help='Use visdom for loss visualization')
+parser.add_argument('--send_images_to_visdom', type=str2bool, default=False, help='Sample a random image from every 10th batch, send it to visdom after augmentations step')
 parser.add_argument('--save_folder', default='weights/', help='Directory for saving checkpoint models')
 parser.add_argument('--dataset_root', default=COCO_ROOT, help='Dataset root directory path')
 parser.add_argument('-f', default=None, type=str, help="Dummy arg so we can load in Jupyter Notebooks")
@@ -104,63 +104,41 @@ def weights_init(m):
 def train():
     net.train()
     # loss counters
-    loc_loss = 0  # epoch
+    loc_loss = 0
     conf_loss = 0
     epoch = 0
     print('Loading Dataset...')
     dataset = COCODetection(args.dataset_root, args.image_set, SSDAugmentation(
         SSD_DIM, MEANS), COCOAnnotationTransform())
 
     epoch_size = len(dataset) // args.batch_size
-    print('Training SSD on ', dataset.name)
+    print('Training SSD on', dataset.name)
     step_index = 0
+
     if args.visdom:
-        # initialize visdom loss plot
-        lot = viz.line(
-            X=torch.zeros((1,)).cpu(),
-            Y=torch.zeros((1, 3)).cpu(),
-            opts=dict(
-                xlabel='Iteration',
-                ylabel='Loss',
-                title='Current SSD Training Loss',
-                legend=['Loc Loss', 'Conf Loss', 'Loss']
-            )
-        )
-        epoch_lot = viz.line(
-            X=torch.zeros((1,)).cpu(),
-            Y=torch.zeros((1, 3)).cpu(),
-            opts=dict(
-                xlabel='Epoch',
-                ylabel='Loss',
-                title='Epoch SSD Training Loss',
-                legend=['Loc Loss', 'Conf Loss', 'Loss']
-            )
-        )
-    batch_iterator = None
+        vis_title = 'SSD.PyTorch on ' + args.image_set
+        vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
+        iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
+        epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
     data_loader = data.DataLoader(dataset, args.batch_size,
                                   num_workers=args.num_workers,
                                   shuffle=True, collate_fn=detection_collate,
                                   pin_memory=True)
+    # create batch iterator
+    batch_iterator = iter(data_loader)
     for iteration in range(args.start_iter, args.max_iter):
-        if (not batch_iterator) or (iteration % epoch_size == 0):
-            # create batch iterator
-            batch_iterator = iter(data_loader)
-        if iteration in STEP_VALUES:
-            step_index += 1
-            adjust_learning_rate(optimizer, args.gamma, step_index)
-            if args.visdom:
-                viz.line(
-                    X=torch.ones((1, 3)).cpu() * epoch,
-                    Y=torch.Tensor([loc_loss, conf_loss,
-                        loc_loss + conf_loss]).unsqueeze(0).cpu() / epoch_size,
-                    win=epoch_lot,
-                    update='append'
-                )
+        if iteration != 0 and (iteration % epoch_size == 0) and args.visdom:
+            update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,
+                            'append', epoch_size)
             # reset epoch loss counters
             loc_loss = 0
             conf_loss = 0
             epoch += 1
 
+        if iteration in STEP_VALUES:
+            step_index += 1
+            adjust_learning_rate(optimizer, args.gamma, step_index)
+
         # load train data
         images, targets = next(batch_iterator)
 
@@ -182,29 +160,15 @@ def train():
         t1 = time.time()
         loc_loss += loss_l.data[0]
         conf_loss += loss_c.data[0]
+
         if iteration % 10 == 0:
-            print('Timer: %.4f sec.' % (t1 - t0))
+            print('timer: %.4f sec.' % (t1 - t0))
             print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
-            if args.visdom and args.send_images_to_visdom:
-                random_batch_index = np.random.randint(images.size(0))
-                viz.image(images.data[random_batch_index].cpu().numpy())
+
         if args.visdom:
-            viz.line(
-                X=torch.ones((1, 3)).cpu() * iteration,
-                Y=torch.Tensor([loss_l.data[0], loss_c.data[0],
-                    loss_l.data[0] + loss_c.data[0]]).unsqueeze(0).cpu(),
-                win=lot,
-                update='append'
-            )
-            # hacky fencepost solution for 0th epoch plot
-            if iteration == 0:
-                viz.line(
-                    X=torch.zeros((1, 3)).cpu(),
-                    Y=torch.Tensor([loc_loss, conf_loss,
-                        loc_loss + conf_loss]).unsqueeze(0).cpu(),
-                    win=epoch_lot,
-                    update=True
-                )
+            update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
+                            iter_plot, epoch_plot, 'append')
+
         if iteration % 5000 == 0:
             print('Saving state, iter:', iteration)
             torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +
@@ -224,5 +188,36 @@ def adjust_learning_rate(optimizer, gamma, step):
         param_group['lr'] = lr
 
 
+def create_vis_plot(_xlabel, _ylabel, _title, _legend):
+    return viz.line(
+        X=torch.zeros((1,)).cpu(),
+        Y=torch.zeros((1, 3)).cpu(),
+        opts=dict(
+            xlabel=_xlabel,
+            ylabel=_ylabel,
+            title=_title,
+            legend=_legend
+        )
+    )
+
+
+def update_vis_plot(iteration, loc, conf, window1, window2, update_type,
+                    epoch_size=1):
+    viz.line(
+        X=torch.ones((1, 3)).cpu() * iteration,
+        Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu() / epoch_size,
+        win=window1,
+        update=update_type
+    )
+    # initialize epoch plot on first iteration
+    if iteration == 0:
+        viz.line(
+            X=torch.zeros((1, 3)).cpu(),
+            Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu(),
+            win=window2,
+            update=True
+        )
+
+
 if __name__ == '__main__':
     train()