refactor DataContainer and datasets

mmdet/datasets/__init__.py

 from .coco import CocoDataset
-from .collate import *
-from .sampler import *
-from .transforms import *
+
+__all__ = ['CocoDataset']

mmdet/datasets/coco.py

@@ -7,7 +7,7 @@ from torch.utils.data import Dataset
 
 from .transforms import (ImageTransform, BboxTransform, PolyMaskTransform,
                          Numpy2Tensor)
-from .utils import show_ann, random_scale
+from .utils import to_tensor, show_ann, random_scale
 from .utils import DataContainer as DC
 
 
@@ -71,6 +71,7 @@ def parse_ann_info(ann_info, cat2label, with_mask=True):
 
 
 class CocoDataset(Dataset):
+
     def __init__(self,
                  ann_file,
                  img_prefix,
@@ -227,27 +228,28 @@ class CocoDataset(Dataset):
                         ann['mask_polys'], ann['poly_lens'],
                         img_info['height'], img_info['width'], flip)
 
-            ori_shape = (img_info['height'], img_info['width'])
+            ori_shape = (img_info['height'], img_info['width'], 3)
             img_meta = dict(
-                ori_shape=DC(ori_shape),
-                img_shape=DC(img_shape),
-                scale_factor=DC(scale_factor),
-                flip=DC(flip))
+                ori_shape=ori_shape,
+                img_shape=img_shape,
+                scale_factor=scale_factor,
+                flip=flip)
 
             data = dict(
-                img=DC(img, stack=True),
-                img_meta=img_meta,
-                gt_bboxes=DC(gt_bboxes))
+                img=DC(to_tensor(img), stack=True),
+                img_meta=DC(img_meta, cpu_only=True),
+                gt_bboxes=DC(to_tensor(gt_bboxes)))
             if self.proposals is not None:
-                data['proposals'] = DC(proposals)
+                data['proposals'] = DC(to_tensor(proposals))
             if self.with_label:
-                data['gt_labels'] = DC(gt_labels)
+                data['gt_labels'] = DC(to_tensor(gt_labels))
             if self.with_crowd:
-                data['gt_bboxes_ignore'] = DC(gt_bboxes_ignore)
+                data['gt_bboxes_ignore'] = DC(to_tensor(gt_bboxes_ignore))
             if self.with_mask:
-                data['gt_mask_polys'] = DC(gt_mask_polys)
-                data['gt_poly_lens'] = DC(gt_poly_lens)
-                data['num_polys_per_mask'] = DC(num_polys_per_mask)
+                data['gt_masks'] = dict(
+                    polys=DC(gt_mask_polys, cpu_only=True),
+                    poly_lens=DC(gt_poly_lens, cpu_only=True),
+                    polys_per_mask=DC(num_polys_per_mask, cpu_only=True))
             return data
 
     def prepare_test_img(self, idx):
@@ -258,37 +260,37 @@ class CocoDataset(Dataset):
                     if self.proposals is not None else None)
 
         def prepare_single(img, scale, flip, proposal=None):
-            _img, _img_shape, _scale_factor = self.img_transform(
+            _img, img_shape, scale_factor = self.img_transform(
                 img, scale, flip)
-            img, img_shape, scale_factor = self.numpy2tensor(
-                _img, _img_shape, _scale_factor)
-            ori_shape = (img_info['height'], img_info['width'])
-            img_meta = dict(
-                ori_shape=ori_shape,
+            _img = to_tensor(_img)
+            _img_meta = dict(
+                ori_shape=(img_info['height'], img_info['width'], 3),
                 img_shape=img_shape,
                 scale_factor=scale_factor,
                 flip=flip)
             if proposal is not None:
-                proposal = self.bbox_transform(proposal, _scale_factor, flip)
-                proposal = self.numpy2tensor(proposal)
-            return img, img_meta, proposal
+                _proposal = self.bbox_transform(proposal, scale_factor, flip)
+                _proposal = to_tensor(_proposal)
+            else:
+                _proposal = None
+            return _img, _img_meta, _proposal
 
         imgs = []
         img_metas = []
         proposals = []
         for scale in self.img_scales:
-            img, img_meta, proposal = prepare_single(img, scale, False,
-                                                     proposal)
-            imgs.append(img)
-            img_metas.append(img_meta)
-            proposals.append(proposal)
+            _img, _img_meta, _proposal = prepare_single(
+                img, scale, False, proposal)
+            imgs.append(_img)
+            img_metas.append(DC(_img_meta, cpu_only=True))
+            proposals.append(_proposal)
             if self.flip_ratio > 0:
-                img, img_meta, prop = prepare_single(img, scale, True,
-                                                     proposal)
-                imgs.append(img)
-                img_metas.append(img_meta)
-                proposals.append(prop)
-        if self.proposals is None:
-            return imgs, img_metas
-        else:
-            return imgs, img_metas, proposals
+                _img, _img_meta, _proposal = prepare_single(
+                    img, scale, True, proposal)
+                imgs.append(_img)
+                img_metas.append(DC(_img_meta, cpu_only=True))
+                proposals.append(_proposal)
+        data = dict(img=imgs, img_meta=img_metas)
+        if self.proposals is not None:
+            data['proposals'] = proposals
+        return data

mmdet/datasets/data_engine.py

-from functools import partial
-import torch
-from .coco import CocoDataset
-from .collate import collate
-from .sampler import GroupSampler, DistributedGroupSampler
-
-
-def build_data(cfg, args):
-    dataset = CocoDataset(**cfg)
-
-    if args.dist:
-        sampler = DistributedGroupSampler(dataset, args.img_per_gpu,
-                                     args.world_size, args.rank)
-        batch_size = args.img_per_gpu
-        num_workers = args.data_workers
-    else:
-        sampler = GroupSampler(dataset, args.img_per_gpu)
-        batch_size = args.world_size * args.img_per_gpu
-        num_workers = args.world_size * args.data_workers
-
-    loader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=args.img_per_gpu,
-        sampler=sampler,
-        num_workers=num_workers,
-        collate_fn=partial(collate, samples_per_gpu=args.img_per_gpu),
-        pin_memory=False)
-
-    return loader

mmdet/datasets/loader/__init__.py

+from .build_loader import build_dataloader
+from .collate import collate
+from .sampler import GroupSampler, DistributedGroupSampler
+
+__all__ = [
+    'collate', 'GroupSampler', 'DistributedGroupSampler', 'build_dataloader'
+]

mmdet/datasets/loader/build_loader.py

+from functools import partial
+
+from torch.utils.data import DataLoader
+
+from .collate import collate
+from .sampler import GroupSampler, DistributedGroupSampler
+
+
+def build_dataloader(dataset,
+                     imgs_per_gpu,
+                     workers_per_gpu,
+                     num_gpus,
+                     dist=True,
+                     world_size=1,
+                     rank=0,
+                     **kwargs):
+    if dist:
+        sampler = DistributedGroupSampler(dataset, imgs_per_gpu, world_size,
+                                          rank)
+        batch_size = imgs_per_gpu
+        num_workers = workers_per_gpu
+    else:
+        sampler = GroupSampler(dataset, imgs_per_gpu)
+        batch_size = num_gpus * imgs_per_gpu
+        num_workers = num_gpus * workers_per_gpu
+
+    if not kwargs.get('shuffle', True):
+        sampler = None
+
+    data_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        num_workers=num_workers,
+        collate_fn=partial(collate, samples_per_gpu=imgs_per_gpu),
+        pin_memory=False,
+        **kwargs)
+
+    return data_loader

mmdet/datasets/collate.py → mmdet/datasets/loader/collate.py

@@ -4,17 +4,24 @@ import torch
 import torch.nn.functional as F
 from torch.utils.data.dataloader import default_collate
 
-from .utils import DataContainer
+from ..utils import DataContainer
 
 # https://github.com/pytorch/pytorch/issues/973
 import resource
 rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
 resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))
 
-__all__ = ['collate']
-
 
 def collate(batch, samples_per_gpu=1):
+    """Puts each data field into a tensor/DataContainer with outer dimension
+    batch size.
+
+    Extend default_collate to add support for :type:`~mmdet.DataContainer`.
+    There are 3 cases for data containers.
+    1. cpu_only = True, e.g., meta data
+    2. cpu_only = False, stack = True, e.g., images tensors
+    3. cpu_only = False, stack = False, e.g., gt bboxes
+    """
 
     if not isinstance(batch, collections.Sequence):
         raise TypeError("{} is not supported.".format(batch.dtype))
@@ -22,7 +29,13 @@ def collate(batch, samples_per_gpu=1):
     if isinstance(batch[0], DataContainer):
         assert len(batch) % samples_per_gpu == 0
         stacked = []
-        if batch[0].stack:
+        if batch[0].cpu_only:
+            for i in range(0, len(batch), samples_per_gpu):
+                stacked.append(
+                    [sample.data for sample in batch[i:i + samples_per_gpu]])
+            return DataContainer(
+                stacked, batch[0].stack, batch[0].padding_value, cpu_only=True)
+        elif batch[0].stack:
             for i in range(0, len(batch), samples_per_gpu):
                 assert isinstance(batch[i].data, torch.Tensor)
                 # TODO: handle tensors other than 3d

mmdet/datasets/sampler.py → mmdet/datasets/loader/sampler.py

@@ -7,8 +7,6 @@ import numpy as np
 from torch.distributed import get_world_size, get_rank
 from torch.utils.data.sampler import Sampler
 
-__all__ = ['GroupSampler', 'DistributedGroupSampler']
-
 
 class GroupSampler(Sampler):
 

mmdet/datasets/transforms.py

@@ -29,7 +29,7 @@ class ImageTransform(object):
         self.size_divisor = size_divisor
 
     def __call__(self, img, scale, flip=False):
-        img, scale_factor = mmcv.imrescale(img, scale, True)
+        img, scale_factor = mmcv.imrescale(img, scale, return_scale=True)
         img_shape = img.shape
         img = mmcv.imnorm(img, self.mean, self.std, self.to_rgb)
         if flip:
@@ -39,76 +39,20 @@ class ImageTransform(object):
         img = img.transpose(2, 0, 1)
         return img, img_shape, scale_factor
 
-        # img, scale = cvb.resize_keep_ar(img_or_path, max_long_edge,
-        #                                 max_short_edge, True)
-        # shape_scale = np.array(img.shape + (scale, ), dtype=np.float32)
-        # if flip:
-        #     img = img[:, ::-1, :].copy()
-        # if self.color_order == 'RGB':
-        #     img = cvb.bgr2rgb(img)
-        # img = img.astype(np.float32)
-        # img -= self.color_mean
-        # img /= self.color_std
-        # if self.size_divisor is None:
-        #     padded_img = img
-        # else:
-        #     pad_h = int(np.ceil(
-        #         img.shape[0] / self.size_divisor)) * self.size_divisor
-        #     pad_w = int(np.ceil(
-        #         img.shape[1] / self.size_divisor)) * self.size_divisor
-        #     padded_img = cvb.pad_img(img, (pad_h, pad_w), pad_val=0)
-        # padded_img = padded_img.transpose(2, 0, 1)
-        # return padded_img, shape_scale
-
-
-class ImageCrop(object):
-    """crop image patches and resize patches into fixed size
-    1. (read and) flip image (if needed)
-    2. crop image patches according to given bboxes
-    3. resize patches into fixed size (default 224x224)
-    4. normalize the image (if needed)
-    5. transpose to (c, h, w) (if needed)
-    """
 
-    def __init__(self,
-                 normalize=True,
-                 transpose=True,
-                 color_order='RGB',
-                 color_mean=(0, 0, 0),
-                 color_std=(1, 1, 1)):
-        self.normalize = normalize
-        self.transpose = transpose
-
-        assert color_order in ['RGB', 'BGR']
-        self.color_order = color_order
-        self.color_mean = np.array(color_mean, dtype=np.float32)
-        self.color_std = np.array(color_std, dtype=np.float32)
-
-    def __call__(self,
-                 img_or_path,
-                 bboxes,
-                 crop_size,
-                 scale_ratio=1.0,
-                 flip=False):
-        img = cvb.read_img(img_or_path)
-        if flip:
-            img = img[:, ::-1, :].copy()
-        crop_imgs = cvb.crop_img(
-            img,
-            bboxes[:, :4],
-            scale_ratio=scale_ratio,
-            pad_fill=self.color_mean)
-        processed_crop_imgs_list = []
-        for i in range(len(crop_imgs)):
-            crop_img = crop_imgs[i]
-            crop_img = cvb.resize(crop_img, crop_size)
-            crop_img = crop_img.astype(np.float32)
-            crop_img -= self.color_mean
-            crop_img /= self.color_std
-            processed_crop_imgs_list.append(crop_img)
-        processed_crop_imgs = np.stack(processed_crop_imgs_list, axis=0)
-        processed_crop_imgs = processed_crop_imgs.transpose(0, 3, 1, 2)
-        return processed_crop_imgs
+def bbox_flip(bboxes, img_shape):
+    """Flip bboxes horizontally.
+
+    Args:
+        bboxes(ndarray): shape (..., 4*k)
+        img_shape(tuple): (height, width)
+    """
+    assert bboxes.shape[-1] % 4 == 0
+    w = img_shape[1]
+    flipped = bboxes.copy()
+    flipped[..., 0::4] = w - bboxes[..., 2::4] - 1
+    flipped[..., 2::4] = w - bboxes[..., 0::4] - 1
+    return flipped
 
 
 class BboxTransform(object):
@@ -124,7 +68,7 @@ class BboxTransform(object):
     def __call__(self, bboxes, img_shape, scale_factor, flip=False):
         gt_bboxes = bboxes * scale_factor
         if flip:
-            gt_bboxes = mmcv.bbox_flip(gt_bboxes, img_shape)
+            gt_bboxes = bbox_flip(gt_bboxes, img_shape)
         gt_bboxes[:, 0::2] = np.clip(gt_bboxes[:, 0::2], 0, img_shape[1])
         gt_bboxes[:, 1::2] = np.clip(gt_bboxes[:, 1::2], 0, img_shape[0])
         if self.max_num_gts is None:
@@ -161,42 +105,6 @@ class PolyMaskTransform(object):
         return gt_mask_polys, gt_poly_lens, num_polys_per_mask
 
 
-class MaskTransform(object):
-    """Preprocess masks
-    1. resize masks to expected size and stack to a single array
-    2. flip the masks (if needed)
-    3. pad the masks (if needed)
-    """
-
-    def __init__(self, max_num_gts, pad_size=None):
-        self.max_num_gts = max_num_gts
-        self.pad_size = pad_size
-
-    def __call__(self, masks, img_size, flip=False):
-        max_long_edge = max(img_size)
-        max_short_edge = min(img_size)
-        masks = [
-            cvb.resize_keep_ar(
-                mask,
-                max_long_edge,
-                max_short_edge,
-                interpolation=cvb.INTER_NEAREST) for mask in masks
-        ]
-        masks = np.stack(masks, axis=0)
-        if flip:
-            masks = masks[:, ::-1, :]
-        if self.pad_size is None:
-            pad_h = masks.shape[1]
-            pad_w = masks.shape[2]
-        else:
-            pad_size = self.pad_size if self.pad_size > 0 else max_long_edge
-            pad_h = pad_w = pad_size
-        padded_masks = np.zeros(
-            (self.max_num_gts, pad_h, pad_w), dtype=masks.dtype)
-        padded_masks[:masks.shape[0], :masks.shape[1], :masks.shape[2]] = masks
-        return padded_masks
-
-
 class Numpy2Tensor(object):
 
     def __init__(self):

mmdet/datasets/utils/data_container.py

 import functools
-from collections import Sequence
 
-import mmcv
-import numpy as np
 import torch
 
 
-def to_tensor(data):
-    """Convert objects of various python types to :obj:`torch.Tensor`.
-
-    Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
-    :class:`Sequence`, :class:`int` and :class:`float`.
-    """
-    if isinstance(data, np.ndarray):
-        return torch.from_numpy(data)
-    elif isinstance(data, torch.Tensor):
-        return data
-    elif isinstance(data, Sequence) and not mmcv.is_str(data):
-        return torch.tensor(data)
-    elif isinstance(data, int):
-        return torch.LongTensor([data])
-    elif isinstance(data, float):
-        return torch.FloatTensor([data])
-    else:
-        raise TypeError('type {} cannot be converted to tensor.'.format(
-            type(data)))
-
-
 def assert_tensor_type(func):
 
     @functools.wraps(func)
@@ -41,11 +17,9 @@ def assert_tensor_type(func):
 
 class DataContainer(object):
 
-    def __init__(self, data, stack=False, padding_value=0):
-        if isinstance(data, list):
-            self._data = data
-        else:
-            self._data = to_tensor(data)
+    def __init__(self, data, stack=False, padding_value=0, cpu_only=False):
+        self._data = data
+        self._cpu_only = cpu_only
         self._stack = stack
         self._padding_value = padding_value
 
@@ -63,6 +37,10 @@ class DataContainer(object):
         else:
             return type(self.data)
 
+    @property
+    def cpu_only(self):
+        return self._cpu_only
+
     @property
     def stack(self):
         return self._stack

mmdet/datasets/utils/misc.py

+from collections import Sequence
+
 import mmcv
+import torch
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pycocotools.mask as maskUtils
 
 
+def to_tensor(data):
+    """Convert objects of various python types to :obj:`torch.Tensor`.
+
+    Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
+    :class:`Sequence`, :class:`int` and :class:`float`.
+    """
+    if isinstance(data, torch.Tensor):
+        return data
+    elif isinstance(data, np.ndarray):
+        return torch.from_numpy(data)
+    elif isinstance(data, Sequence) and not mmcv.is_str(data):
+        return torch.tensor(data)
+    elif isinstance(data, int):
+        return torch.LongTensor([data])
+    elif isinstance(data, float):
+        return torch.FloatTensor([data])
+    else:
+        raise TypeError('type {} cannot be converted to tensor.'.format(
+            type(data)))
+
+
 def random_scale(img_scales, mode='range'):
     """Randomly select a scale from a list of scales or scale ranges.