Repository: open-mmlab/mmpose

Issue:
'ConcatDataset' is not available in MMPOSE
My config is something like this. I want to use the 'ConcatDataset' just like it in the MMDETECTION and MMSEGMENT. However, the 'ConcatDataset' is not available in MMPOSE
```
# myconfig.py
data = dict(
    train=dict(
         type = 'ConcatDataset',
         datasets = [dict(
             type='BottomUpCocoDataset',
             ann_file='data/keypoints/2021-11-02-bwrat_side6-kp_val.json',
             img_prefix='data/keypoints/2021-11-02-bwrat_side6-kp/',
             data_cfg=data_cfg,
             pipeline=test_pipeline,
            dataset_info=dataset_info),
         dict(
             type='BottomUpCocoDataset',
             ann_file='data/keypoints/2021-11-24-wrat_side6-kp_val.json',
             img_prefix='data/keypoints/2021-11-24-wrat_side6-kp/',
             data_cfg=data_cfg,
             pipeline=val_pipeline,
             dataset_info=dataset_info)
         ])
)
```


Discussion:
Good suggestion! 
Currently ConcatDataset is not supported in MMPose. We plan to add this feature soon.
我曾尝试移植MMDETECTIO的 ConcatDataset 到 `mmpose/datasets/dataset_wrappers.py` 和`mmpose/datasets/builder.py`。但后续遇到其它报错。具体来说就是在train> evaluation interval 时会报错。
```python
#mmpose/datasets/dataset_wrappers.py
@DATASETS.register_module()
class ConcatDataset(_ConcatDataset):
    """A wrapper of concatenated dataset.

    Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but
    concat the group flag for image aspect ratio.

    Args:
        datasets (list[:obj:`Dataset`]): A list of datasets.
        separate_eval (bool): Whether to evaluate the results
            separately if it is used as validation dataset.
            Defaults to True.
    """

    def __init__(self, datasets, separate_eval=True):
        super(ConcatDataset, self).__init__(datasets)
        # self.CLASSES = datasets[0].CLASSES
        self.separate_eval = separate_eval
        if not separate_eval:
            if len(set([type(ds) for ds in datasets])) != 1:
                raise NotImplementedError(
                    'All the datasets should have same types')

        if hasattr(datasets[0], 'flag'):
            flags = []
            for i in range(0, len(datasets)):
                flags.append(datasets[i].flag)
            self.flag = np.concatenate(flags)


    def evaluate(self, results, logger=None, **kwargs):
        """Evaluate the results.

        Args:
            results (list[list | tuple]): Testing results of the dataset.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.

        Returns:
            dict[str: float]: AP results of the total dataset or each separate
            dataset if `self.separate_eval=True`.
        """
        assert len(results) == self.cumulative_sizes[-1], \
            ('Dataset and results have different sizes: '
             f'{self.cumulative_sizes[-1]} v.s. {len(results)}')

        # Check whether all the datasets support evaluation
        for dataset in self.datasets:
            assert hasattr(dataset, 'evaluate'), \
                    f'{type(dataset)} does not implement evaluate function'

        if self.separate_eval:
            dataset_idx = -1
            total_eval_results = dict()
            for size, dataset in zip(self.cumulative_sizes, self.datasets):
                start_idx = 0 if dataset_idx == -1 else \
                    self.cumulative_sizes[dataset_idx]
                end_idx = self.cumulative_sizes[dataset_idx + 1]

                results_per_dataset = results[start_idx:end_idx]
                print_log(
                    f'\nEvaluateing {dataset.ann_file} with '
                    f'{len(results_per_dataset)} images now',
                    logger=logger)

                eval_results_per_dataset = dataset.evaluate(
                    results_per_dataset, logger=logger, **kwargs)
                dataset_idx += 1
                for k, v in eval_results_per_dataset.items():
                    total_eval_results.update({f'{dataset_idx}_{k}': v})

            return total_eval_results
        elif len(set([type(ds) for ds in self.datasets])) != 1:
            raise NotImplementedError(
                'All the datasets should have same types')
        else:
            original_data_infos = self.datasets[0].data_infos
            self.datasets[0].data_infos = sum(
                [dataset.data_infos for dataset in self.datasets], [])
            eval_results = self.datasets[0].evaluate(
                results, logger=logger, **kwargs)
            self.datasets[0].data_infos = original_data_infos
            return eval_results
```

```python
# mmpose/datasets/builder.py
def build_dataset(cfg, default_args=None):
    """Build a dataset from config dict.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        default_args (dict, optional): Default initialization arguments.
            Default: None.

    Returns:
        Dataset: The constructed dataset.
    """
    from .dataset_wrappers import RepeatDataset
    from .dataset_wrappers import ConcatDataset

    if cfg['type'] == 'RepeatDataset':
        dataset = RepeatDataset(
            build_dataset(cfg['dataset'], default_args), cfg['times'])
    elif isinstance(cfg, (list, tuple)):
            dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg])
    elif cfg['type'] == 'ConcatDataset':
        dataset = ConcatDataset(
            [build_dataset(c, default_args) for c in cfg['datasets']],
            cfg.get('separate_eval', True))
    else:
        dataset = build_from_cfg(cfg, DATASETS, default_args)
    return dataset
```

train> evaluation interval  过程的bug
```
2021-12-27 15:04:33,308 - mmpose - INFO - Saving checkpoint at 10 epochs
[>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>] 452/452, 9.3 task/s, elapsed: 49s, ETA:     0s2021-12-27 15:05:22,968 - mmpose - INFO - 
Evaluateing data/rat_keypoints/2021-11-02-bwrat_side6-kp_val.json with 179 images now
/home/liying_lab/chenxinfeng/.conda/envs/mmpose/lib/python3.7/site-packages/json_tricks/encoders.py:367: UserWarning: json-tricks: numpy scalar serialization is experimental and may work differently in future versions
  warnings.warn('json-tricks: numpy scalar serialization is experimental and may work differently in future versions')
Loading and preparing results...
DONE (t=0.00s)
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *keypoints*
DONE (t=0.05s).
Accumulating evaluation results...
DONE (t=0.00s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] =  0.626
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets= 20 ] =  0.967
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets= 20 ] =  0.721
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] = -1.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] =  0.626
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] =  0.675
 Average Recall     (AR) @[ IoU=0.50      | area=   all | maxDets= 20 ] =  0.972
 Average Recall     (AR) @[ IoU=0.75      | area=   all | maxDets= 20 ] =  0.777
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] = -1.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] =  0.675
2021-12-27 15:05:23,194 - mmpose - INFO - 
Evaluateing data/rat_keypoints/2021-11-24-wrat_side6-kp_val.json with 273 images now
Loading and preparing results...
DONE (t=0.01s)
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *keypoints*
DONE (t=0.07s).
Accumulating evaluation results...
DONE (t=0.00s).
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] =  0.729
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets= 20 ] =  0.950
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets= 20 ] =  0.856
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] =  0.000
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] =  0.737
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 20 ] =  0.775
 Average Recall     (AR) @[ IoU=0.50      | area=   all | maxDets= 20 ] =  0.956
 Average Recall     (AR) @[ IoU=0.75      | area=   all | maxDets= 20 ] =  0.890
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets= 20 ] =  0.000
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets= 20 ] =  0.784
Traceback (most recent call last):
  File "tools/train.py", line 170, in <module>
    main()
  File "tools/train.py", line 166, in main
    meta=meta)
  File "/DATA/chenxinfeng/mmpose/mmpose/apis/train.py", line 190, in train_model
    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
  File "/home/liying_lab/chenxinfeng/.conda/envs/mmpose/lib/python3.7/site-packages/mmcv/runner/epoch_based_runner.py", line 127, in run
    epoch_runner(data_loaders[i], **kwargs)
  File "/home/liying_lab/chenxinfeng/.conda/envs/mmpose/lib/python3.7/site-packages/mmcv/runner/epoch_based_runner.py", line 54, in train
    self.call_hook('after_train_epoch')
  File "/home/liying_lab/chenxinfeng/.conda/envs/mmpose/lib/python3.7/site-packages/mmcv/runner/base_runner.py", line 307, in call_hook
    getattr(hook, fn_name)(self)
  File "/home/liying_lab/chenxinfeng/.conda/envs/mmpose/lib/python3.7/site-packages/mmcv/runner/hooks/evaluation.py", line 267, in after_train_epoch
    self._do_evaluate(runner)
  File "/home/liying_lab/chenxinfeng/.conda/envs/mmpose/lib/python3.7/site-packages/mmcv/runner/hooks/evaluation.py", line 273, in _do_evaluate
    key_score = self.evaluate(runner, results)
  File "/DATA/chenxinfeng/mmpose/mmpose/core/evaluation/eval_hooks.py", line 80, in evaluate
    return eval_res[self.key_indicator]
KeyError: 'AP'
```
@Canwang-sjtu  Could you please take a look at this issue?

Resolving patch:
diff --git a/mmpose/datasets/builder.py b/mmpose/datasets/builder.py
--- a/mmpose/datasets/builder.py
+++ b/mmpose/datasets/builder.py
@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import copy
 import platform
 import random
 from functools import partial
@@ -6,8 +7,9 @@
 import numpy as np
 from mmcv.parallel import collate
 from mmcv.runner import get_dist_info
-from mmcv.utils import Registry, build_from_cfg
+from mmcv.utils import Registry, build_from_cfg, is_seq_of
 from mmcv.utils.parrots_wrapper import _get_dataloader
+from torch.utils.data.dataset import ConcatDataset
 
 from .samplers import DistributedSampler
 
@@ -24,6 +26,39 @@
 PIPELINES = Registry('pipeline')
 
 
+def _concat_dataset(cfg, default_args=None):
+    types = cfg['type']
+    ann_files = cfg['ann_file']
+    img_prefixes = cfg.get('img_prefix', None)
+    dataset_infos = cfg.get('dataset_info', None)
+
+    num_joints = cfg['data_cfg'].get('num_joints', None)
+    dataset_channel = cfg['data_cfg'].get('dataset_channel', None)
+
+    datasets = []
+    num_dset = len(ann_files)
+    for i in range(num_dset):
+        cfg_copy = copy.deepcopy(cfg)
+        cfg_copy['ann_file'] = ann_files[i]
+
+        if isinstance(types, (list, tuple)):
+            cfg_copy['type'] = types[i]
+        if isinstance(img_prefixes, (list, tuple)):
+            cfg_copy['img_prefix'] = img_prefixes[i]
+        if isinstance(dataset_infos, (list, tuple)):
+            cfg_copy['dataset_info'] = dataset_infos[i]
+
+        if isinstance(num_joints, (list, tuple)):
+            cfg_copy['data_cfg']['num_joints'] = num_joints[i]
+
+        if is_seq_of(dataset_channel, list):
+            cfg_copy['data_cfg']['dataset_channel'] = dataset_channel[i]
+
+        datasets.append(build_dataset(cfg_copy, default_args))
+
+    return ConcatDataset(datasets)
+
+
 def build_dataset(cfg, default_args=None):
     """Build a dataset from config dict.
 
@@ -37,9 +72,16 @@ def build_dataset(cfg, default_args=None):
     """
     from .dataset_wrappers import RepeatDataset
 
-    if cfg['type'] == 'RepeatDataset':
+    if isinstance(cfg, (list, tuple)):
+        dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg])
+    elif cfg['type'] == 'ConcatDataset':
+        dataset = ConcatDataset(
+            [build_dataset(c, default_args) for c in cfg['datasets']])
+    elif cfg['type'] == 'RepeatDataset':
         dataset = RepeatDataset(
             build_dataset(cfg['dataset'], default_args), cfg['times'])
+    elif isinstance(cfg.get('ann_file'), (list, tuple)):
+        dataset = _concat_dataset(cfg, default_args)
     else:
         dataset = build_from_cfg(cfg, DATASETS, default_args)
     return dataset


Test patch:
diff --git a/tests/test_datasets/test_dataset_wrapper.py b/tests/test_datasets/test_dataset_wrapper.py
new file mode 100644
--- /dev/null
+++ b/tests/test_datasets/test_dataset_wrapper.py
@@ -0,0 +1,67 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmcv import Config
+
+from mmpose.datasets.builder import build_dataset
+
+
+def test_concat_dataset():
+    # build COCO-like dataset config
+    dataset_info = Config.fromfile(
+        'configs/_base_/datasets/coco.py').dataset_info
+
+    channel_cfg = dict(
+        num_output_channels=17,
+        dataset_joints=17,
+        dataset_channel=[
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
+        ],
+        inference_channel=[
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
+        ])
+
+    data_cfg = dict(
+        image_size=[192, 256],
+        heatmap_size=[48, 64],
+        num_output_channels=channel_cfg['num_output_channels'],
+        num_joints=channel_cfg['dataset_joints'],
+        dataset_channel=channel_cfg['dataset_channel'],
+        inference_channel=channel_cfg['inference_channel'],
+        soft_nms=False,
+        nms_thr=1.0,
+        oks_thr=0.9,
+        vis_thr=0.2,
+        use_gt_bbox=True,
+        det_bbox_thr=0.0,
+        bbox_file='tests/data/coco/test_coco_det_AP_H_56.json',
+    )
+
+    dataset_cfg = dict(
+        type='TopDownCocoDataset',
+        ann_file='tests/data/coco/test_coco.json',
+        img_prefix='tests/data/coco/',
+        data_cfg=data_cfg,
+        pipeline=[],
+        dataset_info=dataset_info)
+
+    dataset = build_dataset(dataset_cfg)
+
+    # Case 1: build ConcatDataset explicitly
+    concat_dataset_cfg = dict(
+        type='ConcatDataset', datasets=[dataset_cfg, dataset_cfg])
+    concat_dataset = build_dataset(concat_dataset_cfg)
+    assert len(concat_dataset) == 2 * len(dataset)
+
+    # Case 2: build ConcatDataset from cfg sequence
+    concat_dataset = build_dataset([dataset_cfg, dataset_cfg])
+    assert len(concat_dataset) == 2 * len(dataset)
+
+    # Case 3: build ConcatDataset from ann_file sequence
+    concat_dataset_cfg = dataset_cfg.copy()
+    for key in ['ann_file', 'type', 'img_prefix', 'dataset_info']:
+        val = concat_dataset_cfg[key]
+        concat_dataset_cfg[key] = [val] * 2
+    for key in ['num_joints', 'dataset_channel']:
+        val = concat_dataset_cfg['data_cfg'][key]
+        concat_dataset_cfg['data_cfg'][key] = [val] * 2
+    concat_dataset = build_dataset(concat_dataset_cfg)
+    assert len(concat_dataset) == 2 * len(dataset)