| |
| import copy |
|
|
| from mmdet.datasets import DATASETS, ConcatDataset, build_dataset |
|
|
| from mmocr.utils import is_2dlist, is_type_list |
|
|
|
|
| @DATASETS.register_module() |
| class UniformConcatDataset(ConcatDataset): |
| """A wrapper of ConcatDataset which support dataset pipeline assignment and |
| replacement. |
| |
| Args: |
| datasets (list[dict] | list[list[dict]]): A list of datasets cfgs. |
| separate_eval (bool): Whether to evaluate the results |
| separately if it is used as validation dataset. |
| Defaults to True. |
| pipeline (None | list[dict] | list[list[dict]]): If ``None``, |
| each dataset in datasets use its own pipeline; |
| If ``list[dict]``, it will be assigned to the dataset whose |
| pipeline is None in datasets; |
| If ``list[list[dict]]``, pipeline of dataset which is None |
| in datasets will be replaced by the corresponding pipeline |
| in the list. |
| force_apply (bool): If True, apply pipeline above to each dataset |
| even if it have its own pipeline. Default: False. |
| """ |
|
|
| def __init__(self, |
| datasets, |
| separate_eval=True, |
| pipeline=None, |
| force_apply=False, |
| **kwargs): |
| new_datasets = [] |
| if pipeline is not None: |
| assert isinstance( |
| pipeline, |
| list), 'pipeline must be list[dict] or list[list[dict]].' |
| if is_type_list(pipeline, dict): |
| self._apply_pipeline(datasets, pipeline, force_apply) |
| new_datasets = datasets |
| elif is_2dlist(pipeline): |
| assert is_2dlist(datasets) |
| assert len(datasets) == len(pipeline) |
| for sub_datasets, tmp_pipeline in zip(datasets, pipeline): |
| self._apply_pipeline(sub_datasets, tmp_pipeline, |
| force_apply) |
| new_datasets.extend(sub_datasets) |
| else: |
| if is_2dlist(datasets): |
| for sub_datasets in datasets: |
| new_datasets.extend(sub_datasets) |
| else: |
| new_datasets = datasets |
| datasets = [build_dataset(c, kwargs) for c in new_datasets] |
| super().__init__(datasets, separate_eval) |
|
|
| @staticmethod |
| def _apply_pipeline(datasets, pipeline, force_apply=False): |
| from_cfg = all(isinstance(x, dict) for x in datasets) |
| assert from_cfg, 'datasets should be config dicts' |
| assert all(isinstance(x, dict) for x in pipeline) |
| for dataset in datasets: |
| if dataset['pipeline'] is None or force_apply: |
| dataset['pipeline'] = copy.deepcopy(pipeline) |
|
|
