fix size_index bug (issue longcw#72)

Author: longcw

README.md

@@ -81,14 +81,10 @@ and set the path in `yolo2-pytorch/cfgs/exps/darknet19_exp1.py`.
 
 7. (optional) Training with TensorBoard.
 
-    To use the TensorBoard, install Crayon (https://github.com/torrvision/crayon)
-    How to use the crayon
-    ```
-    docker pull alband/crayon
-    docker run -d -p 8888:8888 -p 8889:8889 --name crayon alband/crayon
-    ```
-
-and set `use_tensorboard = True` in `yolo2-pytorch/cfgs/config.py`.
+    To use the TensorBoard, 
+    set `use_tensorboard = True` in `yolo2-pytorch/cfgs/config.py`
+    and install TensorboardX (https://github.com/lanpa/tensorboard-pytorch).
+    Tensorboard log will be saved in `training/runs`.
 
 
 6. Run the training program: `python train.py`.

datasets/imdb.py

@@ -15,12 +15,6 @@ def mkdir(path, max_depth=3):
         os.mkdir(path)
 
 
-def image_resize(im, size_index):
-    w, h = cfg.multi_scale_inp_size[size_index]
-    im = cv2.resize(im, (w, h))
-    return im
-
-
 class ImageDataset(object):
     def __init__(self, name, datadir, batch_size, im_processor,
                  processes=3, shuffle=True, dst_size=None):
@@ -51,22 +45,42 @@ def next_batch(self, size_index):
         batch = {'images': [], 'gt_boxes': [], 'gt_classes': [],
                  'dontcare': [], 'origin_im': []}
         i = 0
+        if self.gen is None:
+            indexes = np.arange(len(self.image_names), dtype=np.int)
+            if self._shuffle:
+                np.random.shuffle(indexes)
+            self.gen = self.pool.imap(partial(self._im_processor,
+                                              size_index=None),
+                                      ([self.image_names[i],
+                                        self.get_annotation(i),
+                                        self.dst_size] for i in indexes),
+                                      chunksize=self.batch_size)
+            self._epoch += 1
+            print(('epoch {} start...'.format(self._epoch)))
+
         while i < self.batch_size:
             try:
                 images, gt_boxes, classes, dontcare, origin_im = next(self.gen)
-                images = image_resize(images, size_index)
+
+                # multi-scale
+                w, h = cfg.multi_scale_inp_size[size_index]
+                gt_boxes = gt_boxes.astype(float)
+                gt_boxes[:, 0::2] *= float(w) / images.shape[1]
+                gt_boxes[:, 1::2] *= float(h) / images.shape[0]
+                images = cv2.resize(images, (w, h))
+
                 batch['images'].append(images)
                 batch['gt_boxes'].append(gt_boxes)
                 batch['gt_classes'].append(classes)
                 batch['dontcare'].append(dontcare)
                 batch['origin_im'].append(origin_im)
                 i += 1
-            except (StopIteration, AttributeError, TypeError):
+            except (StopIteration,):
                 indexes = np.arange(len(self.image_names), dtype=np.int)
                 if self._shuffle:
                     np.random.shuffle(indexes)
                 self.gen = self.pool.imap(partial(self._im_processor,
-                                                  size_index=size_index),
+                                                  size_index=None),
                                           ([self.image_names[i],
                                             self.get_annotation(i),
                                             self.dst_size] for i in indexes),

demo.py

@@ -9,6 +9,11 @@
 from utils.timer import Timer
 import cfgs.config as cfg
 
+# This prevents deadlocks in the data loader, caused by
+# some incompatibility between pytorch and cv2 multiprocessing.
+# See https://github.com/pytorch/pytorch/issues/1355.
+cv2.setNumThreads(0)
+
 
 def preprocess(fname):
     # return fname

test.py

@@ -12,13 +12,6 @@
 import cfgs.config as cfg
 
 
-def preprocess(fname):
-    # return fname
-    image = cv2.imread(fname)
-    im_data = np.expand_dims(yolo_utils.preprocess_test(image, cfg.multi_scale_inp_size), 0)  # noqa
-    return image, im_data
-
-
 parser = argparse.ArgumentParser(description='PyTorch Yolo')
 parser.add_argument('--image_size_index', type=int, default=0,
                     metavar='image_size_index',

train.py

@@ -12,9 +12,9 @@
 from random import randint
 
 try:
-    from pycrayon import CrayonClient
+    from tensorboardX import SummaryWriter
 except ImportError:
-    CrayonClient = None
+    SummaryWriter = None
 
 
 # data loader
@@ -42,22 +42,12 @@
                             weight_decay=cfg.weight_decay)
 
 # tensorboad
-use_tensorboard = cfg.use_tensorboard and CrayonClient is not None
+use_tensorboard = cfg.use_tensorboard and SummaryWriter is not None
 # use_tensorboard = False
-remove_all_log = False
 if use_tensorboard:
-    cc = CrayonClient(hostname='127.0.0.1')
-    if remove_all_log:
-        print('remove all experiments')
-        cc.remove_all_experiments()
-    if start_epoch == 0:
-        try:
-            cc.remove_experiment(cfg.exp_name)
-        except ValueError:
-            pass
-        exp = cc.create_experiment(cfg.exp_name)
-    else:
-        exp = cc.open_experiment(cfg.exp_name)
+    summary_writer = SummaryWriter(os.path.join(cfg.TRAIN_DIR, 'runs', cfg.exp_name))
+else:
+    summary_writer = None
 
 batch_per_epoch = imdb.batch_per_epoch
 train_loss = 0
@@ -81,7 +71,7 @@
     im_data = net_utils.np_to_variable(im,
                                        is_cuda=True,
                                        volatile=False).permute(0, 3, 1, 2)
-    net(im_data, gt_boxes, gt_classes, dontcare, size_index)
+    bbox_pred, iou_pred, prob_pred = net(im_data, gt_boxes, gt_classes, dontcare, size_index)
 
     # backward
     loss = net.loss
@@ -106,12 +96,22 @@
                 iou_loss, cls_loss, duration,
                 str(datetime.timedelta(seconds=int((batch_per_epoch - step_cnt) * duration))))))  # noqa
 
-        if use_tensorboard and step % cfg.log_interval == 0:
-            exp.add_scalar_value('loss_train', train_loss, step=step)
-            exp.add_scalar_value('loss_bbox', bbox_loss, step=step)
-            exp.add_scalar_value('loss_iou', iou_loss, step=step)
-            exp.add_scalar_value('loss_cls', cls_loss, step=step)
-            exp.add_scalar_value('learning_rate', lr, step=step)
+        if summary_writer and step % cfg.log_interval == 0:
+            summary_writer.add_scalar('loss_train', train_loss, step)
+            summary_writer.add_scalar('loss_bbox', bbox_loss, step)
+            summary_writer.add_scalar('loss_iou', iou_loss, step)
+            summary_writer.add_scalar('loss_cls', cls_loss, step)
+            summary_writer.add_scalar('learning_rate', lr, step)
+
+            # plot results
+            bbox_pred = bbox_pred.data[0:1].cpu().numpy()
+            iou_pred = iou_pred.data[0:1].cpu().numpy()
+            prob_pred = prob_pred.data[0:1].cpu().numpy()
+            image = im[0]
+            bboxes, scores, cls_inds = yolo_utils.postprocess(
+                bbox_pred, iou_pred, prob_pred, image.shape, cfg, thresh=0.3, size_index=size_index)
+            im2show = yolo_utils.draw_detection(image, bboxes, scores, cls_inds, cfg)
+            summary_writer.add_image('predict', im2show, step)
 
         train_loss = 0
         bbox_loss, iou_loss, cls_loss = 0., 0., 0.

utils/yolo.py

@@ -6,6 +6,12 @@
 from utils.cython_yolo import yolo_to_bbox
 
 
+# This prevents deadlocks in the data loader, caused by
+# some incompatibility between pytorch and cv2 multiprocessing.
+# See https://github.com/pytorch/pytorch/issues/1355.
+cv2.setNumThreads(0)
+
+
 def clip_boxes(boxes, im_shape):
     """
     Clip boxes to image boundaries.
@@ -50,7 +56,7 @@ def _offset_boxes(boxes, im_shape, scale, offs, flip):
 
 def preprocess_train(data, size_index):
     im_path, blob, inp_size = data
-    inp_size = inp_size[size_index]
+
     boxes, gt_classes = blob['boxes'], blob['gt_classes']
 
     im = cv2.imread(im_path)
@@ -60,7 +66,8 @@ def preprocess_train(data, size_index):
     scale, offs, flip = trans_param
     boxes = _offset_boxes(boxes, im.shape, scale, offs, flip)
 
-    if inp_size is not None:
+    if inp_size is not None and size_index is not None:
+        inp_size = inp_size[size_index]
         w, h = inp_size
         boxes[:, 0::2] *= float(w) / im.shape[1]
         boxes[:, 1::2] *= float(h) / im.shape[0]
@@ -126,7 +133,7 @@ def postprocess(bbox_pred, iou_pred, prob_pred, im_shape, cfg, thresh=0.05,
     prob_pred = prob_pred[(np.arange(prob_pred.shape[0]), cls_inds)]
     scores = iou_pred * prob_pred
     # scores = iou_pred
-
+    assert len(scores) == len(bbox_pred), '{}, {}'.format(scores.shape, bbox_pred.shape)
     # threshold
     keep = np.where(scores >= thresh)
     bbox_pred = bbox_pred[keep]