support recursion

mmdet/datasets/__init__.py

@@ -7,6 +7,7 @@ from .repeat_dataset import RepeatDataset
 
 __all__ = [
     'CustomDataset', 'CocoDataset', 'GroupSampler', 'DistributedGroupSampler',
-    'ConcatDataset', 'build_dataloader', 'to_tensor', 'random_scale',
-    'show_ann', 'get_dataset', 'RepeatDataset'
+    'build_dataloader', 'to_tensor', 'random_scale', 'show_ann',
+    'get_dataset', 'ExtraAugmentation', 'ConcatDataset', 'RepeatDataset',
 ]
+

mmdet/datasets/repeat_dataset.py

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

mmdet/datasets/utils.py

@@ -75,10 +75,8 @@ def show_ann(coco, img, ann_info):
 
 
 def get_dataset(data_cfg):
-    repeat_times = None
     if data_cfg['type'] == 'RepeatDataset':
-        repeat_times = data_cfg['repeat_times']
-        data_cfg = data_cfg['dataset']
+        return RepeatDataset(get_dataset(data_cfg['type']), data_cfg['times'])
 
     if isinstance(data_cfg['ann_file'], (list, tuple)):
         ann_files = data_cfg['ann_file']
@@ -114,7 +112,4 @@ def get_dataset(data_cfg):
         dset = ConcatDataset(dsets)
     else:
         dset = dsets[0]
-
-    if repeat_times is not None:
-        dset = RepeatDataset(dset, repeat_times)
     return dset