97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import copy
|
|
import os.path as osp
|
|
from typing import List
|
|
|
|
from mmengine.fileio import load
|
|
|
|
from mmseg.registry import DATASETS
|
|
from .basesegdataset import BaseSegDataset
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class DecathlonDataset(BaseSegDataset):
|
|
"""Dataset for Dacathlon dataset.
|
|
|
|
The dataset.json format is shown as follows
|
|
|
|
.. code-block:: none
|
|
|
|
{
|
|
"name": "BRATS",
|
|
"tensorImageSize": "4D",
|
|
"modality":
|
|
{
|
|
"0": "FLAIR",
|
|
"1": "T1w",
|
|
"2": "t1gd",
|
|
"3": "T2w"
|
|
},
|
|
"labels": {
|
|
"0": "background",
|
|
"1": "edema",
|
|
"2": "non-enhancing tumor",
|
|
"3": "enhancing tumour"
|
|
},
|
|
"numTraining": 484,
|
|
"numTest": 266,
|
|
"training":
|
|
[
|
|
{
|
|
"image": "./imagesTr/BRATS_306.nii.gz"
|
|
"label": "./labelsTr/BRATS_306.nii.gz"
|
|
...
|
|
}
|
|
]
|
|
"test":
|
|
[
|
|
"./imagesTs/BRATS_557.nii.gz"
|
|
...
|
|
]
|
|
}
|
|
"""
|
|
|
|
def load_data_list(self) -> List[dict]:
|
|
"""Load annotation from directory or annotation file.
|
|
|
|
Returns:
|
|
list[dict]: All data info of dataset.
|
|
"""
|
|
# `self.ann_file` denotes the absolute annotation file path if
|
|
# `self.root=None` or relative path if `self.root=/path/to/data/`.
|
|
annotations = load(self.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)}!')
|
|
raw_data_list = annotations[
|
|
'training'] if not self.test_mode else annotations['test']
|
|
data_list = []
|
|
for raw_data_info in raw_data_list:
|
|
# `2:` works for removing './' in file path, which will break
|
|
# loading from cloud storage.
|
|
if isinstance(raw_data_info, dict):
|
|
data_info = dict(
|
|
img_path=osp.join(self.data_root, raw_data_info['image']
|
|
[2:]))
|
|
data_info['seg_map_path'] = osp.join(
|
|
self.data_root, raw_data_info['label'][2:])
|
|
else:
|
|
data_info = dict(
|
|
img_path=osp.join(self.data_root, raw_data_info)[2:])
|
|
data_info['label_map'] = self.label_map
|
|
data_info['reduce_zero_label'] = self.reduce_zero_label
|
|
data_info['seg_fields'] = []
|
|
data_list.append(data_info)
|
|
annotations.pop('training')
|
|
annotations.pop('test')
|
|
|
|
metainfo = copy.deepcopy(annotations)
|
|
metainfo['classes'] = [*metainfo['labels'].values()]
|
|
# Meta information load from annotation file will not influence the
|
|
# existed meta information load from `BaseDataset.METAINFO` and
|
|
# `metainfo` arguments defined in constructor.
|
|
for k, v in metainfo.items():
|
|
self._metainfo.setdefault(k, v)
|
|
|
|
return data_list
|