Merge pull request #144 from yhcao6/RepeatDataset

repeat dataset

mmdet/datasets/__init__.py

@@ -3,9 +3,10 @@ from .coco import CocoDataset
 from .loader import GroupSampler, DistributedGroupSampler, build_dataloader
 from .utils import to_tensor, random_scale, show_ann, get_dataset
 from .concat_dataset import ConcatDataset
+from .repeat_dataset import RepeatDataset
 
 __all__ = [
     'CustomDataset', 'CocoDataset', 'GroupSampler', 'DistributedGroupSampler',
-    'ConcatDataset', 'build_dataloader', 'to_tensor', 'random_scale',
-    'show_ann', 'get_dataset'
+    'build_dataloader', 'to_tensor', 'random_scale', 'show_ann',
+    'get_dataset', 'ExtraAugmentation', 'ConcatDataset', 'RepeatDataset',
 ]

mmdet/datasets/repeat_dataset.py

+import numpy as np
+
+
+class RepeatDataset(object):
+
+    def __init__(self, dataset, times):
+        self.dataset = dataset
+        self.times = times
+        if hasattr(self.dataset, 'flag'):
+            self.flag = np.tile(self.dataset.flag, times)
+        self._original_length = len(self.dataset)
+
+    def __getitem__(self, idx):
+        return self.dataset[idx % self._original_length]
+
+    def __len__(self):
+        return self.times * self._original_length

mmdet/datasets/utils.py

@@ -8,6 +8,7 @@ import torch
 import matplotlib.pyplot as plt
 import numpy as np
 from .concat_dataset import ConcatDataset
+from .repeat_dataset import RepeatDataset
 from .. import datasets
 
 
@@ -74,6 +75,10 @@ def show_ann(coco, img, ann_info):
 
 
 def get_dataset(data_cfg):
+    if data_cfg['type'] == 'RepeatDataset':
+        return RepeatDataset(
+            get_dataset(data_cfg['dataset']), data_cfg['times'])
+
     if isinstance(data_cfg['ann_file'], (list, tuple)):
         ann_files = data_cfg['ann_file']
         num_dset = len(ann_files)