| | |
| | import collections |
| | import copy |
| | from typing import List, Optional, Sequence, Union |
| |
|
| | from mmengine.dataset import ConcatDataset, force_full_init |
| |
|
| | from mmseg.registry import DATASETS, TRANSFORMS |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class MultiImageMixDataset: |
| | """A wrapper of multiple images mixed dataset. |
| | |
| | Suitable for training on multiple images mixed data augmentation like |
| | mosaic and mixup. |
| | |
| | Args: |
| | dataset (ConcatDataset or dict): The dataset to be mixed. |
| | pipeline (Sequence[dict]): Sequence of transform object or |
| | config dict to be composed. |
| | skip_type_keys (list[str], optional): Sequence of type string to |
| | be skip pipeline. Default to None. |
| | """ |
| |
|
| | def __init__(self, |
| | dataset: Union[ConcatDataset, dict], |
| | pipeline: Sequence[dict], |
| | skip_type_keys: Optional[List[str]] = None, |
| | lazy_init: bool = False) -> None: |
| | assert isinstance(pipeline, collections.abc.Sequence) |
| |
|
| | if isinstance(dataset, dict): |
| | self.dataset = DATASETS.build(dataset) |
| | elif isinstance(dataset, ConcatDataset): |
| | self.dataset = dataset |
| | else: |
| | raise TypeError( |
| | 'elements in datasets sequence should be config or ' |
| | f'`ConcatDataset` instance, but got {type(dataset)}') |
| |
|
| | if skip_type_keys is not None: |
| | assert all([ |
| | isinstance(skip_type_key, str) |
| | for skip_type_key in skip_type_keys |
| | ]) |
| | self._skip_type_keys = skip_type_keys |
| |
|
| | self.pipeline = [] |
| | self.pipeline_types = [] |
| | for transform in pipeline: |
| | if isinstance(transform, dict): |
| | self.pipeline_types.append(transform['type']) |
| | transform = TRANSFORMS.build(transform) |
| | self.pipeline.append(transform) |
| | else: |
| | raise TypeError('pipeline must be a dict') |
| |
|
| | self._metainfo = self.dataset.metainfo |
| | self.num_samples = len(self.dataset) |
| |
|
| | self._fully_initialized = False |
| | if not lazy_init: |
| | self.full_init() |
| |
|
| | @property |
| | def metainfo(self) -> dict: |
| | """Get the meta information of the multi-image-mixed dataset. |
| | |
| | Returns: |
| | dict: The meta information of multi-image-mixed dataset. |
| | """ |
| | return copy.deepcopy(self._metainfo) |
| |
|
| | def full_init(self): |
| | """Loop to ``full_init`` each dataset.""" |
| | if self._fully_initialized: |
| | return |
| |
|
| | self.dataset.full_init() |
| | self._ori_len = len(self.dataset) |
| | self._fully_initialized = True |
| |
|
| | @force_full_init |
| | def get_data_info(self, idx: int) -> dict: |
| | """Get annotation by index. |
| | |
| | Args: |
| | idx (int): Global index of ``ConcatDataset``. |
| | |
| | Returns: |
| | dict: The idx-th annotation of the datasets. |
| | """ |
| | return self.dataset.get_data_info(idx) |
| |
|
| | @force_full_init |
| | def __len__(self): |
| | return self.num_samples |
| |
|
| | def __getitem__(self, idx): |
| | results = copy.deepcopy(self.dataset[idx]) |
| | for (transform, transform_type) in zip(self.pipeline, |
| | self.pipeline_types): |
| | if self._skip_type_keys is not None and \ |
| | transform_type in self._skip_type_keys: |
| | continue |
| |
|
| | if hasattr(transform, 'get_indices'): |
| | indices = transform.get_indices(self.dataset) |
| | if not isinstance(indices, collections.abc.Sequence): |
| | indices = [indices] |
| | mix_results = [ |
| | copy.deepcopy(self.dataset[index]) for index in indices |
| | ] |
| | results['mix_results'] = mix_results |
| |
|
| | results = transform(results) |
| |
|
| | if 'mix_results' in results: |
| | results.pop('mix_results') |
| |
|
| | return results |
| |
|
| | def update_skip_type_keys(self, skip_type_keys): |
| | """Update skip_type_keys. |
| | |
| | It is called by an external hook. |
| | |
| | Args: |
| | skip_type_keys (list[str], optional): Sequence of type |
| | string to be skip pipeline. |
| | """ |
| | assert all([ |
| | isinstance(skip_type_key, str) for skip_type_key in skip_type_keys |
| | ]) |
| | self._skip_type_keys = skip_type_keys |
| |
|
