Merge pull request #2 from OceanPang/dev

padded mask during maskrcnn test

mmdet/models/backbones/resnet.py

@@ -200,7 +200,8 @@ class ResNet(nn.Module):
                  frozen_stages=-1,
                  style='fb',
                  sync_bn=False,
-                 with_cp=False):
+                 with_cp=False,
+                 strict_frozen=False):
         super(ResNet, self).__init__()
         if not len(layers) == len(strides) == len(dilations):
             raise ValueError(
@@ -241,6 +242,8 @@ class ResNet(nn.Module):
         self.feat_dim = block.expansion * 64 * 2**(len(layers) - 1)
         self.with_cp = with_cp
 
+        self.strict_frozen = strict_frozen
+
     def init_weights(self, pretrained=None):
         if isinstance(pretrained, str):
             logger = logging.getLogger()
@@ -278,6 +281,9 @@ class ResNet(nn.Module):
             for m in self.modules():
                 if isinstance(m, nn.BatchNorm2d):
                     m.eval()
+                    if self.strict_frozen:
+                        for params in m.parameters():
+                            params.requires_grad = False
         if mode and self.frozen_stages >= 0:
             for param in self.conv1.parameters():
                 param.requires_grad = False
@@ -310,7 +316,8 @@ def resnet(depth,
            frozen_stages=-1,
            style='fb',
            sync_bn=False,
-           with_cp=False):
+           with_cp=False,
+           strict_frozen=False):
     """Constructs a ResNet model.
 
     Args:
@@ -324,5 +331,5 @@ def resnet(depth,
         raise KeyError('invalid depth {} for resnet'.format(depth))
     block, layers = resnet_cfg[depth]
     model = ResNet(block, layers[:num_stages], strides, dilations, out_indices,
-                   frozen_stages, style, sync_bn, with_cp)
+                   frozen_stages, style, sync_bn, with_cp, strict_frozen)
     return model

mmdet/models/mask_heads/fcn_mask_head.py

@@ -87,7 +87,7 @@ class FCNMaskHead(nn.Module):
         return mask_pred
 
     def get_mask_target(self, pos_proposals, pos_assigned_gt_inds, gt_masks,
-                     img_meta, rcnn_train_cfg):
+                        img_meta, rcnn_train_cfg):
         mask_targets = mask_target(pos_proposals, pos_assigned_gt_inds,
                                    gt_masks, img_meta, rcnn_train_cfg)
         return mask_targets
@@ -117,12 +117,19 @@ class FCNMaskHead(nn.Module):
         if isinstance(mask_pred, torch.Tensor):
             mask_pred = mask_pred.sigmoid().cpu().numpy()
         assert isinstance(mask_pred, np.ndarray)
+
         cls_segms = [[] for _ in range(self.num_classes - 1)]
+        mask_size = mask_pred.shape[-1]
         bboxes = det_bboxes.cpu().numpy()[:, :4]
         labels = det_labels.cpu().numpy() + 1
         img_h = ori_shape[0]
         img_w = ori_shape[1]
 
+        scale = (mask_size + 2.0) / mask_size
+        bboxes = np.round(self._bbox_scaling(bboxes, scale)).astype(np.int32)
+        padded_mask = np.zeros(
+            (mask_size + 2, mask_size + 2), dtype=np.float32)
+
         for i in range(bboxes.shape[0]):
             bbox = bboxes[i, :].astype(int)
             label = labels[i]
@@ -132,18 +139,45 @@ class FCNMaskHead(nn.Module):
             h = max(h, 1)
 
             if not self.class_agnostic:
-                mask_pred_ = mask_pred[i, label, :, :]
+                padded_mask[1:-1, 1:-1] = mask_pred[i, label, :, :]
             else:
-                mask_pred_ = mask_pred[i, 0, :, :]
+                padded_mask[1:-1, 1:-1] = mask_pred[i, 0, :, :]
+            mask = mmcv.imresize(padded_mask, (w, h))
+            mask = np.array(
+                mask > rcnn_test_cfg.mask_thr_binary, dtype=np.uint8)
+            im_mask = np.zeros((img_h, img_w), dtype=np.uint8)
+
+            x0 = max(bbox[0], 0)
+            x1 = min(bbox[2] + 1, img_w)
+            y0 = max(bbox[1], 0)
+            y1 = min(bbox[3] + 1, img_h)
 
-            im_mask = np.zeros((img_h, img_w), dtype=np.float32)
+            im_mask[y0:y1, x0:x1] = mask[(y0 - bbox[1]):(y1 - bbox[1]), (
+                x0 - bbox[0]):(x1 - bbox[0])]
 
-            im_mask[bbox[1]:bbox[1] + h, bbox[0]:bbox[0] + w] = mmcv.imresize(
-                mask_pred_, (w, h))
-            # im_mask = cv2.resize(im_mask, (img_w, img_h))
-            im_mask = np.array(
-                im_mask > rcnn_test_cfg.mask_thr_binary, dtype=np.uint8)
             rle = mask_util.encode(
                 np.array(im_mask[:, :, np.newaxis], order='F'))[0]
             cls_segms[label - 1].append(rle)
         return cls_segms
+
+    def _bbox_scaling(self, bboxes, scale, clip_shape=None):
+        """Scaling bboxes and clip the boundary(optional)
+        Args:
+            bboxes(ndarray): shape(..., 4)
+            scale(float): scaling factor
+            clip(None or tuple): (h, w)
+        Returns:
+            ndarray: scaled bboxes
+        """
+        if float(scale) == 1.0:
+            scaled_bboxes = bboxes.copy()
+        else:
+            w = bboxes[..., 2] - bboxes[..., 0] + 1
+            h = bboxes[..., 3] - bboxes[..., 1] + 1
+            dw = (w * (scale - 1)) * 0.5
+            dh = (h * (scale - 1)) * 0.5
+            scaled_bboxes = bboxes + np.stack((-dw, -dh, dw, dh), axis=-1)
+        if clip_shape is not None:
+            return bbox_clip(scaled_bboxes, clip_shape)
+        else:
+            return scaled_bboxes

tools/train.py

@@ -116,11 +116,6 @@ def main():
         build_dataloader(train_dataset, cfg.data.imgs_per_gpu,
                          cfg.data.workers_per_gpu, cfg.gpus, dist)
     ]
-    if args.validate:
-        val_dataset = obj_from_dict(cfg.data.val, datasets)
-        data_loaders.append(
-            build_dataloader(val_dataset, cfg.data.imgs_per_gpu,
-                             cfg.data.workers_per_gpu, cfg.gpus, dist))
 
     # build model
     model = build_detector(
@@ -133,6 +128,7 @@ def main():
     # build runner
     runner = Runner(model, batch_processor, cfg.optimizer, cfg.work_dir,
                     cfg.log_level)
+
     # register hooks
     optimizer_config = DistOptimizerHook(
         **cfg.optimizer_config) if dist else cfg.optimizer_config
@@ -141,10 +137,11 @@ def main():
     if dist:
         runner.register_hook(DistSamplerSeedHook())
         # register eval hooks
-        if isinstance(model.module, RPN):
-            runner.register_hook(CocoDistEvalRecallHook(cfg.data.val))
-        elif cfg.data.val.type == 'CocoDataset':
-            runner.register_hook(CocoDistEvalmAPHook(cfg.data.val))
+        if args.validate:
+            if isinstance(model.module, RPN):
+                runner.register_hook(CocoDistEvalRecallHook(cfg.data.val))
+            elif cfg.data.val.type == 'CocoDataset':
+                runner.register_hook(CocoDistEvalmAPHook(cfg.data.val))
 
     if cfg.resume_from:
         runner.resume(cfg.resume_from)