# Copyright (c) OpenMMLab. All rights reserved. import copy import os.path as osp from os import PathLike from typing import Optional, Sequence import mmengine from mmcv.transforms import Compose from mmengine.fileio import get_file_backend from .builder import DATASETS def expanduser(path): if isinstance(path, (str, PathLike)): return osp.expanduser(path) else: return path def isabs(uri): return osp.isabs(uri) or ('://' in uri) @DATASETS.register_module() class MultiTaskDataset: """Custom dataset for multi-task dataset. To use the dataset, please generate and provide an annotation file in the below format: .. code-block:: json { "metainfo": { "tasks": [ 'gender' 'wear' ] }, "data_list": [ { "img_path": "a.jpg", gt_label:{ "gender": 0, "wear": [1, 0, 1, 0] } }, { "img_path": "b.jpg", gt_label:{ "gender": 1, "wear": [1, 0, 1, 0] } } ] } Assume we put our dataset in the ``data/mydataset`` folder in the repository and organize it as the below format: :: mmpretrain/ └── data └── mydataset ├── annotation │   ├── train.json │   ├── test.json │   └── val.json ├── train │   ├── a.jpg │   └── ... ├── test │   ├── b.jpg │   └── ... └── val ├── c.jpg └── ... We can use the below config to build datasets: .. code:: python >>> from mmpretrain.datasets import build_dataset >>> train_cfg = dict( ... type="MultiTaskDataset", ... ann_file="annotation/train.json", ... data_root="data/mydataset", ... # The `img_path` field in the train annotation file is relative ... # to the `train` folder. ... data_prefix='train', ... ) >>> train_dataset = build_dataset(train_cfg) Or we can put all files in the same folder: :: mmpretrain/ └── data └── mydataset ├── train.json ├── test.json ├── val.json ├── a.jpg ├── b.jpg ├── c.jpg └── ... And we can use the below config to build datasets: .. code:: python >>> from mmpretrain.datasets import build_dataset >>> train_cfg = dict( ... type="MultiTaskDataset", ... ann_file="train.json", ... data_root="data/mydataset", ... # the `data_prefix` is not required since all paths are ... # relative to the `data_root`. ... ) >>> train_dataset = build_dataset(train_cfg) Args: ann_file (str): The annotation file path. It can be either absolute path or relative path to the ``data_root``. metainfo (dict, optional): The extra meta information. It should be a dict with the same format as the ``"metainfo"`` field in the annotation file. Defaults to None. data_root (str, optional): The root path of the data directory. It's the prefix of the ``data_prefix`` and the ``ann_file``. And it can be a remote path like "s3://openmmlab/xxx/". Defaults to None. data_prefix (str, optional): The base folder relative to the ``data_root`` for the ``"img_path"`` field in the annotation file. Defaults to None. pipeline (Sequence[dict]): A list of dict, where each element represents a operation defined in :mod:`mmpretrain.datasets.pipelines`. Defaults to an empty tuple. test_mode (bool): in train mode or test mode. Defaults to False. """ METAINFO = dict() def __init__(self, ann_file: str, metainfo: Optional[dict] = None, data_root: Optional[str] = None, data_prefix: Optional[str] = None, pipeline: Sequence = (), test_mode: bool = False): self.data_root = expanduser(data_root) # Inference the file client if self.data_root is not None: self.file_backend = get_file_backend(uri=self.data_root) else: self.file_backend = None self.ann_file = self._join_root(expanduser(ann_file)) self.data_prefix = self._join_root(data_prefix) self.test_mode = test_mode self.pipeline = Compose(pipeline) self.data_list = self.load_data_list(self.ann_file, metainfo) def _join_root(self, path): """Join ``self.data_root`` with the specified path. If the path is an absolute path, just return the path. And if the path is None, return ``self.data_root``. Examples: >>> self.data_root = 'a/b/c' >>> self._join_root('d/e/') 'a/b/c/d/e' >>> self._join_root('https://openmmlab.com') 'https://openmmlab.com' >>> self._join_root(None) 'a/b/c' """ if path is None: return self.data_root if isabs(path): return path joined_path = self.file_backend.join_path(self.data_root, path) return joined_path @classmethod def _get_meta_info(cls, in_metainfo: dict = None) -> dict: """Collect meta information from the dictionary of meta. Args: in_metainfo (dict): Meta information dict. Returns: dict: Parsed meta information. """ # `cls.METAINFO` will be overwritten by in_meta metainfo = copy.deepcopy(cls.METAINFO) if in_metainfo is None: return metainfo metainfo.update(in_metainfo) return metainfo def load_data_list(self, ann_file, metainfo_override=None): """Load annotations from an annotation file. Args: ann_file (str): Absolute annotation file path if ``self.root=None`` or relative path if ``self.root=/path/to/data/``. Returns: list[dict]: A list of annotation. """ annotations = mmengine.load(ann_file) if not isinstance(annotations, dict): raise TypeError(f'The annotations loaded from annotation file ' f'should be a dict, but got {type(annotations)}!') if 'data_list' not in annotations: raise ValueError('The annotation file must have the `data_list` ' 'field.') metainfo = annotations.get('metainfo', {}) raw_data_list = annotations['data_list'] # Set meta information. assert isinstance(metainfo, dict), 'The `metainfo` field in the '\ f'annotation file should be a dict, but got {type(metainfo)}' if metainfo_override is not None: assert isinstance(metainfo_override, dict), 'The `metainfo` ' \ f'argument should be a dict, but got {type(metainfo_override)}' metainfo.update(metainfo_override) self._metainfo = self._get_meta_info(metainfo) data_list = [] for i, raw_data in enumerate(raw_data_list): try: data_list.append(self.parse_data_info(raw_data)) except AssertionError as e: raise RuntimeError( f'The format check fails during parse the item {i} of ' f'the annotation file with error: {e}') return data_list def parse_data_info(self, raw_data): """Parse raw annotation to target format. This method will return a dict which contains the data information of a sample. Args: raw_data (dict): Raw data information load from ``ann_file`` Returns: dict: Parsed annotation. """ assert isinstance(raw_data, dict), \ f'The item should be a dict, but got {type(raw_data)}' assert 'img_path' in raw_data, \ "The item doesn't have `img_path` field." data = dict( img_path=self._join_root(raw_data['img_path']), gt_label=raw_data['gt_label'], ) return data @property def metainfo(self) -> dict: """Get meta information of dataset. Returns: dict: meta information collected from ``cls.METAINFO``, annotation file and metainfo argument during instantiation. """ return copy.deepcopy(self._metainfo) def prepare_data(self, idx): """Get data processed by ``self.pipeline``. Args: idx (int): The index of ``data_info``. Returns: Any: Depends on ``self.pipeline``. """ results = copy.deepcopy(self.data_list[idx]) return self.pipeline(results) def __len__(self): """Get the length of the whole dataset. Returns: int: The length of filtered dataset. """ return len(self.data_list) def __getitem__(self, idx): """Get the idx-th image and data information of dataset after ``self.pipeline``. Args: idx (int): The index of of the data. Returns: dict: The idx-th image and data information after ``self.pipeline``. """ return self.prepare_data(idx) def __repr__(self): """Print the basic information of the dataset. Returns: str: Formatted string. """ head = 'Dataset ' + self.__class__.__name__ body = [f'Number of samples: \t{self.__len__()}'] if self.data_root is not None: body.append(f'Root location: \t{self.data_root}') body.append(f'Annotation file: \t{self.ann_file}') if self.data_prefix is not None: body.append(f'Prefix of images: \t{self.data_prefix}') # -------------------- extra repr -------------------- tasks = self.metainfo['tasks'] body.append(f'For {len(tasks)} tasks') for task in tasks: body.append(f' {task} ') # ---------------------------------------------------- if len(self.pipeline.transforms) > 0: body.append('With transforms:') for t in self.pipeline.transforms: body.append(f' {t}') lines = [head] + [' ' * 4 + line for line in body] return '\n'.join(lines)