torch_em.data.datasets.medical.segrap

  1"""The SegRap dataset contains annotations for organ-at-risk (OAR) segmentation in head and neck CT scans
  2of nasopharyngeal carcinoma patients.
  3
  4It comprises the training set of the SegRap2023 challenge (https://segrap2023.grand-challenge.org): 120 patients
  5with a pre-aligned pair of a non-contrast and a contrast-enhanced CT scan, and annotations for 45 OARs.
  6
  7NOTE: The label legend is as follows. Since some of the 45 OARs are nested (e.g. the hippocampi inside the
  8temporal lobes, the cochleae inside the middle ears), the challenge distributes the annotations as a single
  9label volume with 54 disjoint sub-parts, where each overlap of two OARs gets its own id. The ids are
 10(see `SUBPART_NAMES` and `LABEL_IDS`):
 11- 1: Brain, 2: BrainStem, 3: Chiasm, 4/5: TemporalLobe_L/R, 6/7: TemporalLobe_Hippocampus_L/R,
 12  8/9: Hippocampus_L/R, 10/11: Eye_L/R, 12/13: Lens_L/R, 14/15: OpticNerve_L/R, 16/17: MiddleEar_L/R,
 13  18/19: IAC_L/R, 20/21: MiddleEar_TympanicCavity_L/R, 22/23: TympanicCavity_L/R,
 14  24/25: MiddleEar_VestibulSemi_L/R, 26/27: VestibulSemi_L/R, 28/29: Cochlea_L/R,
 15  30/31: MiddleEar_ETbone_L/R, 32/33: ETbone_L/R, 34: Pituitary, 35: OralCavity, 36/37: Mandible_L/R,
 16  38/39: Submandibular_L/R, 40/41: Parotid_L/R, 42/43: Mastoid_L/R, 44/45: TMjoint_L/R, 46: SpinalCord,
 17  47: Esophagus, 48: Larynx, 49: Larynx_Glottic, 50: Larynx_Supraglot, 51: Larynx_PharynxConst,
 18  52: PharynxConst, 53: Thyroid, 54: Trachea
 19`OAR_TO_LABEL_IDS` maps each of the 45 OARs to the ids it consists of. It is taken from the official
 20post-processing code of the challenge (https://github.com/HiLab-git/SegRap2023/blob/main/Tutorial/postprocessing.py).
 21The ids were verified on the data: the label volumes contain the ids 0 to 54.
 22
 23The data is a redistribution of the official challenge data at
 24https://huggingface.co/datasets/YongchengYAO/SegRap23-Lite (CC BY-NC 4.0), which holds the unchanged images
 25and Task001 labels of all 120 training cases, renamed after the case ids. The official release at
 26https://segrap2023.grand-challenge.org/dataset/ requires a signed end user agreement, so please make sure
 27that you are allowed to use the data for your purpose.
 28
 29NOTE: The official release can also be used. Download 'SegRap2023_Training_Set_120cases.zip' and
 30'SegRap2023_Training_Set_120cases_OneHot_Labels.zip' as described on the dataset page and extract them into
 31'<path>', such that '<path>/SegRap2023_Training_Set_120cases/segrap_XXXX/image.nii.gz' (and
 32'image_contrast.nii.gz') and '<path>/SegRap2023_Training_Set_120cases_OneHot_Labels/Task001/segrap_XXXX.nii.gz'
 33exist. This dataset will then use the official data instead of the redistribution.
 34
 35This dataset is from the publication https://doi.org/10.1016/j.media.2024.103447.
 36Please cite it if you use this dataset in your research.
 37"""
 38
 39import os
 40from glob import glob
 41from natsort import natsorted
 42from typing import Union, Tuple, Literal, List
 43
 44from torch.utils.data import Dataset, DataLoader
 45
 46import torch_em
 47
 48from .. import util
 49
 50
 51URLS = {
 52    "ct": "https://huggingface.co/datasets/YongchengYAO/SegRap23-Lite/resolve/main/Images-CT.zip",
 53    "ct_contrast": "https://huggingface.co/datasets/YongchengYAO/SegRap23-Lite/resolve/main/Images-contrastCT.zip",
 54    "labels": "https://huggingface.co/datasets/YongchengYAO/SegRap23-Lite/resolve/main/Masks-Task1.zip",
 55}
 56
 57CHECKSUMS = {
 58    "ct": "1e7b849c5f0296200e5ad9503eeeb9b2a8b3360297078d0ce016c8f8335534ef",
 59    "ct_contrast": "000cc8bf9041c49f7e28ad3e183f7af3f9e701b1842031fb17e21a89db0a8f6a",
 60    "labels": "3441d2bd56ecff4485b99a197b36251b09d2030468840f9399419b517691d297",
 61}
 62
 63FOLDER_NAMES = {"ct": "Images-CT", "ct_contrast": "Images-contrastCT", "labels": "Masks-Task1"}
 64
 65# The file names of the two scans in the official release.
 66OFFICIAL_FILE_NAMES = {"ct": "image.nii.gz", "ct_contrast": "image_contrast.nii.gz"}
 67
 68SUBPART_NAMES = [
 69    "Brain", "BrainStem", "Chiasm", "TemporalLobe_L", "TemporalLobe_R", "TemporalLobe_Hippocampus_L",
 70    "TemporalLobe_Hippocampus_R", "Hippocampus_L", "Hippocampus_R", "Eye_L", "Eye_R", "Lens_L", "Lens_R",
 71    "OpticNerve_L", "OpticNerve_R", "MiddleEar_L", "MiddleEar_R", "IAC_L", "IAC_R",
 72    "MiddleEar_TympanicCavity_L", "MiddleEar_TympanicCavity_R", "TympanicCavity_L", "TympanicCavity_R",
 73    "MiddleEar_VestibulSemi_L", "MiddleEar_VestibulSemi_R", "VestibulSemi_L", "VestibulSemi_R", "Cochlea_L",
 74    "Cochlea_R", "MiddleEar_ETbone_L", "MiddleEar_ETbone_R", "ETbone_L", "ETbone_R", "Pituitary", "OralCavity",
 75    "Mandible_L", "Mandible_R", "Submandibular_L", "Submandibular_R", "Parotid_L", "Parotid_R", "Mastoid_L",
 76    "Mastoid_R", "TMjoint_L", "TMjoint_R", "SpinalCord", "Esophagus", "Larynx", "Larynx_Glottic",
 77    "Larynx_Supraglot", "Larynx_PharynxConst", "PharynxConst", "Thyroid", "Trachea",
 78]
 79
 80LABEL_IDS = {"background": 0, **{name: i + 1 for i, name in enumerate(SUBPART_NAMES)}}
 81
 82OAR_TO_LABEL_IDS = {
 83    "Brain": [1, 2, 3, 4, 5, 6, 7, 8, 9],
 84    "BrainStem": [2],
 85    "Chiasm": [3],
 86    "TemporalLobe_L": [4, 6],
 87    "TemporalLobe_R": [5, 7],
 88    "Hippocampus_L": [8, 6],
 89    "Hippocampus_R": [9, 7],
 90    "Eye_L": [10, 12],
 91    "Eye_R": [11, 13],
 92    "Lens_L": [12],
 93    "Lens_R": [13],
 94    "OpticNerve_L": [14],
 95    "OpticNerve_R": [15],
 96    "MiddleEar_L": [18, 16, 20, 24, 28, 30],
 97    "MiddleEar_R": [19, 17, 21, 25, 29, 31],
 98    "IAC_L": [18],
 99    "IAC_R": [19],
100    "TympanicCavity_L": [22, 20],
101    "TympanicCavity_R": [23, 21],
102    "VestibulSemi_L": [26, 24],
103    "VestibulSemi_R": [27, 25],
104    "Cochlea_L": [28],
105    "Cochlea_R": [29],
106    "ETbone_L": [32, 30],
107    "ETbone_R": [33, 31],
108    "Pituitary": [34],
109    "OralCavity": [35],
110    "Mandible_L": [36],
111    "Mandible_R": [37],
112    "Submandibular_L": [38],
113    "Submandibular_R": [39],
114    "Parotid_L": [40],
115    "Parotid_R": [41],
116    "Mastoid_L": [42],
117    "Mastoid_R": [43],
118    "TMjoint_L": [44],
119    "TMjoint_R": [45],
120    "SpinalCord": [46],
121    "Esophagus": [47],
122    "Larynx": [48, 49, 50, 51],
123    "Larynx_Glottic": [49],
124    "Larynx_Supraglot": [50],
125    "PharynxConst": [51, 52],
126    "Thyroid": [53],
127    "Trachea": [54],
128}
129
130
131def _get_official_case_dirs(path):
132    case_dirs = glob(os.path.join(path, "**", "SegRap2023_Training_Set_120cases", "segrap_*"), recursive=True)
133    return natsorted([p for p in case_dirs if os.path.isdir(p)])
134
135
136def _download_component(path, name, download):
137    data_dir = os.path.join(path, FOLDER_NAMES[name])
138    if os.path.exists(data_dir):
139        return data_dir
140
141    os.makedirs(path, exist_ok=True)
142    zip_path = os.path.join(path, f"{FOLDER_NAMES[name]}.zip")
143    util.download_source(path=zip_path, url=URLS[name], download=download, checksum=CHECKSUMS[name])
144    util.unzip(zip_path=zip_path, dst=path)
145
146    return data_dir
147
148
149def get_segrap_data(
150    path: Union[os.PathLike, str], modality: Literal["ct", "ct_contrast"] = "ct", download: bool = False
151) -> str:
152    """Download the SegRap dataset.
153
154    Args:
155        path: Filepath to a folder where the data is downloaded for further processing.
156        modality: The CT scan to download. Either 'ct' (non-contrast) or 'ct_contrast' (contrast-enhanced).
157        download: Whether to download the data if it is not present.
158
159    Returns:
160        Filepath where the data is stored.
161    """
162    if modality not in OFFICIAL_FILE_NAMES:
163        raise ValueError(f"'{modality}' is not a valid modality. Choose one of {list(OFFICIAL_FILE_NAMES)}.")
164
165    if len(_get_official_case_dirs(path)) > 0:  # The official data was downloaded manually.
166        return path
167
168    _download_component(path, "labels", download)
169    _download_component(path, modality, download)
170
171    return path
172
173
174def get_segrap_paths(
175    path: Union[os.PathLike, str], modality: Literal["ct", "ct_contrast"] = "ct", download: bool = False
176) -> Tuple[List[str], List[str]]:
177    """Get paths to the SegRap data.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        modality: The CT scan to use as input. Either 'ct' (non-contrast) or 'ct_contrast' (contrast-enhanced).
182        download: Whether to download the data if it is not present.
183
184    Returns:
185        List of filepaths for the image data.
186        List of filepaths for the label data.
187    """
188    data_dir = get_segrap_data(path, modality, download)
189
190    case_dirs = _get_official_case_dirs(data_dir)
191    if len(case_dirs) > 0:  # The official layout, with one folder per case and the labels in a separate folder.
192        label_dir = os.path.join(os.path.split(os.path.split(case_dirs[0])[0])[0], "Task001")
193        raw_paths = [os.path.join(p, OFFICIAL_FILE_NAMES[modality]) for p in case_dirs]
194        label_paths = [os.path.join(label_dir, f"{os.path.basename(p)}.nii.gz") for p in case_dirs]
195    else:  # The redistributed layout, with the files named after the case ids.
196        raw_paths = natsorted(glob(os.path.join(data_dir, FOLDER_NAMES[modality], "*.nii.gz")))
197        label_paths = [
198            os.path.join(data_dir, FOLDER_NAMES["labels"], os.path.basename(p)) for p in raw_paths
199        ]
200
201    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths + label_paths)
202
203    return raw_paths, label_paths
204
205
206def get_segrap_dataset(
207    path: Union[os.PathLike, str],
208    patch_shape: Tuple[int, ...],
209    modality: Literal["ct", "ct_contrast"] = "ct",
210    resize_inputs: bool = False,
211    download: bool = False,
212    **kwargs
213) -> Dataset:
214    """Get the SegRap dataset for organ-at-risk segmentation.
215
216    Args:
217        path: Filepath to a folder where the data is downloaded for further processing.
218        patch_shape: The patch shape to use for training.
219        modality: The CT scan to use as input. Either 'ct' (non-contrast) or 'ct_contrast' (contrast-enhanced).
220        resize_inputs: Whether to resize inputs to the desired patch shape.
221        download: Whether to download the data if it is not present.
222        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
223
224    Returns:
225        The segmentation dataset.
226    """
227    raw_paths, label_paths = get_segrap_paths(path, modality, download)
228
229    if resize_inputs:
230        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
231        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
232            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
233        )
234
235    return torch_em.default_segmentation_dataset(
236        raw_paths=raw_paths,
237        raw_key="data",
238        label_paths=label_paths,
239        label_key="data",
240        patch_shape=patch_shape,
241        is_seg_dataset=True,
242        **kwargs
243    )
244
245
246def get_segrap_loader(
247    path: Union[os.PathLike, str],
248    batch_size: int,
249    patch_shape: Tuple[int, ...],
250    modality: Literal["ct", "ct_contrast"] = "ct",
251    resize_inputs: bool = False,
252    download: bool = False,
253    **kwargs
254) -> DataLoader:
255    """Get the SegRap dataloader for organ-at-risk segmentation.
256
257    Args:
258        path: Filepath to a folder where the data is downloaded for further processing.
259        batch_size: The batch size for training.
260        patch_shape: The patch shape to use for training.
261        modality: The CT scan to use as input. Either 'ct' (non-contrast) or 'ct_contrast' (contrast-enhanced).
262        resize_inputs: Whether to resize inputs to the desired patch shape.
263        download: Whether to download the data if it is not present.
264        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
265
266    Returns:
267        The DataLoader.
268    """
269    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
270    dataset = get_segrap_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
271    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)