| | |
| | import os |
| | from typing import Dict, List, Optional, Sequence, Union |
| |
|
| | from mmseg.registry import DATASETS |
| | from .basesegdataset import BaseSegDataset |
| |
|
| | try: |
| | from dsdl.dataset import DSDLDataset |
| | except ImportError: |
| | DSDLDataset = None |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class DSDLSegDataset(BaseSegDataset): |
| | """Dataset for dsdl segmentation. |
| | |
| | Args: |
| | specific_key_path(dict): Path of specific key which can not |
| | be loaded by it's field name. |
| | pre_transform(dict): pre-transform functions before loading. |
| | used_labels(sequence): list of actual used classes in train steps, |
| | this must be subset of class domain. |
| | """ |
| |
|
| | METAINFO = {} |
| |
|
| | def __init__(self, |
| | specific_key_path: Dict = {}, |
| | pre_transform: Dict = {}, |
| | used_labels: Optional[Sequence] = None, |
| | **kwargs) -> None: |
| |
|
| | if DSDLDataset is None: |
| | raise RuntimeError( |
| | 'Package dsdl is not installed. Please run "pip install dsdl".' |
| | ) |
| | self.used_labels = used_labels |
| |
|
| | loc_config = dict(type='LocalFileReader', working_dir='') |
| | if kwargs.get('data_root'): |
| | kwargs['ann_file'] = os.path.join(kwargs['data_root'], |
| | kwargs['ann_file']) |
| | required_fields = ['Image', 'LabelMap'] |
| |
|
| | self.dsdldataset = DSDLDataset( |
| | dsdl_yaml=kwargs['ann_file'], |
| | location_config=loc_config, |
| | required_fields=required_fields, |
| | specific_key_path=specific_key_path, |
| | transform=pre_transform, |
| | ) |
| | BaseSegDataset.__init__(self, **kwargs) |
| |
|
| | def load_data_list(self) -> List[Dict]: |
| | """Load data info from a dsdl yaml file named as ``self.ann_file`` |
| | |
| | Returns: |
| | List[dict]: A list of data list. |
| | """ |
| |
|
| | if self.used_labels: |
| | self._metainfo['classes'] = tuple(self.used_labels) |
| | self.label_map = self.get_label_map(self.used_labels) |
| | else: |
| | self._metainfo['classes'] = tuple(['background'] + |
| | self.dsdldataset.class_names) |
| | data_list = [] |
| |
|
| | for i, data in enumerate(self.dsdldataset): |
| | datainfo = dict( |
| | img_path=os.path.join(self.data_prefix['img_path'], |
| | data['Image'][0].location), |
| | seg_map_path=os.path.join(self.data_prefix['seg_map_path'], |
| | data['LabelMap'][0].location), |
| | label_map=self.label_map, |
| | reduce_zero_label=self.reduce_zero_label, |
| | seg_fields=[], |
| | ) |
| | data_list.append(datainfo) |
| |
|
| | return data_list |
| |
|
| | def get_label_map(self, |
| | new_classes: Optional[Sequence] = None |
| | ) -> Union[Dict, None]: |
| | """Require label mapping. |
| | |
| | The ``label_map`` is a dictionary, its keys are the old label ids and |
| | its values are the new label ids, and is used for changing pixel |
| | labels in load_annotations. If and only if old classes in class_dom |
| | is not equal to new classes in args and nether of them is not |
| | None, `label_map` is not None. |
| | Args: |
| | new_classes (list, tuple, optional): The new classes name from |
| | metainfo. Default to None. |
| | Returns: |
| | dict, optional: The mapping from old classes to new classes. |
| | """ |
| | old_classes = ['background'] + self.dsdldataset.class_names |
| | if (new_classes is not None and old_classes is not None |
| | and list(new_classes) != list(old_classes)): |
| |
|
| | label_map = {} |
| | if not set(new_classes).issubset(old_classes): |
| | raise ValueError( |
| | f'new classes {new_classes} is not a ' |
| | f'subset of classes {old_classes} in class_dom.') |
| | for i, c in enumerate(old_classes): |
| | if c not in new_classes: |
| | label_map[i] = 255 |
| | else: |
| | label_map[i] = new_classes.index(c) |
| | return label_map |
| | else: |
| | return None |
| |
|
