torch_em.data.datasets.medical.structseg

  1"""The StructSeg dataset contains annotations for organ-at-risk (OAR) and gross target volume (GTV)
  2segmentation in head and neck and in chest CT scans of nasopharynx and lung cancer patients.
  3
  4It comprises the training sets of the four tasks of the StructSeg2019 challenge
  5(https://structseg2019.grand-challenge.org), each with 50 annotated CT volumes:
  6- task1 ('HaN_OAR'): 22 organs at risk in head and neck CT scans of nasopharynx cancer patients.
  7- task2 ('Naso_GTV'): the gross target volume of nasopharynx cancer, on the same 50 scans as task1.
  8- task3 ('Thoracic_OAR'): 6 organs at risk in chest CT scans of lung cancer patients.
  9- task4 ('Lung_GTV'): the gross target volume of lung cancer, on the same 50 scans as task3.
 10The 10 test volumes per task are distributed without annotations and are therefore not included here.
 11
 12NOTE: The label legend is as follows (see `LABEL_IDS`):
 13- task1: background: 0, brainstem: 1, left eye: 2, right eye: 3, left lens: 4, right lens: 5,
 14  left optic nerve: 6, right optic nerve: 7, optic chiasm: 8, left temporal lobe: 9,
 15  right temporal lobe: 10, pituitary: 11, left parotid gland: 12, right parotid gland: 13,
 16  left inner ear: 14, right inner ear: 15, left middle ear: 16, right middle ear: 17,
 17  left temporomandibular joint: 18, right temporomandibular joint: 19, spinal cord: 20,
 18  left mandible: 21, right mandible: 22
 19- task2: background: 0, nasopharynx cancer gross target volume: 1
 20- task3: background: 0, left lung: 1, right lung: 2, heart: 3, esophagus: 4, trachea: 5, spinal cord: 6
 21- task4: background: 0, lung cancer gross target volume: 1
 22The id order is taken from https://github.com/zhilothebest/Coronary_Calcium, which transcribes the readme of
 23the official release: task1 lists 'Brain_Stem, Eye_L, Eye_R, Lens_L, Lens_R, Opt_Nerve_L, Opt_Nerve_R,
 24Opt_Chiasma, Temporal_Lobes_L, Temporal_Lobes_R, Pituitary, Parotid_Gland_L, Parotid_Gland_R, Inner_Ear_L,
 25Inner_Ear_R, Mid_Ear_L, Mid_Ear_R, T_M_Joint_L, T_M_Joint_R, Spinal_Cord, Mandible_L, Mandible_R' and task3
 26lists 'Lung_L, Lung_R, Heart, Esophagus, Trachea, Spinal_Cord', each 'corresponding to the label 1 to N in
 27the annotation file'. NOTE: Neither order is the order in which the organs are listed on
 28https://structseg2019.grand-challenge.org/Dataset/, and the per-label statistics published in
 29https://github.com/openmedlab/Awesome-Medical-Dataset/blob/main/resources/StructSeg2019_Task3.md carry the
 30task3 names in the website order instead of the id order, so they call id 3 the spinal cord when it is in
 31fact the heart.
 32
 33The task1 ids were verified on the data: all 50 volumes contain exactly the ids 0 - 22, the mean intensity
 34identifies the bony structures (inner ear 1043 HU, mandible 734 HU, temporomandibular joint 408 HU) and the
 35soft tissue ones (eye 25 HU, brainstem 29 HU, parotid gland 22 HU), the mean volumes match the expected
 36anatomy (lens 0.27 cm3, pituitary 0.58 cm3, optic chiasm 1.05 cm3, eye 9.1 cm3, brainstem 27 cm3,
 37temporal lobe 106 cm3), and in the LPS orientation of the volumes every structure named '_l' lies on the
 38patient's left of its '_r' partner.
 39
 40The task3 ids 1 - 4 were confirmed independently of that transcription, against the copy of the task3 volumes
 41redistributed in https://huggingface.co/datasets/blueyo0/SA-Med3D-140K (its 'ct_general_StructSeg2019_subtask2'
 42entries), which stores one named binary mask per organ. Over the 50 cases the volumes of its 'lung_left',
 43'lung_right', 'heart' and 'esophagus' masks match the statistics published for the ids 1, 2, 3 and 4 to within
 440.3% (median 1324 / 1842 / 690 / 35.7 cm3 against 1325 / 1841 / 692 / 35.5 cm3, and the minima and maxima agree
 45just as closely). The ids 5 and 6 could not be checked that way, since that copy holds neither a trachea nor a
 46spinal cord mask, so they rest on the transcribed readme alone.
 47
 48The data of 'task1' is downloaded from the redistribution at
 49https://huggingface.co/datasets/Luffy503/VoCo_Downstream (the file 'Dataset190_Structseg19.zip'), which is
 50published as part of the VoCo benchmark (https://github.com/Luffy03/Large-Scale-Medical) and is tagged
 51Apache-2.0 by its uploader. NOTE: This is a third-party redistribution of data that the official release
 52only hands out to registered challenge participants, and the StructSeg organizers did not grant a license
 53for it. The participant agreement is quoted in https://github.com/zhilothebest/Coronary_Calcium as saying
 54that 'the datasets can not be publicly posted, or distributed to anyone outside of this project', so please
 55make sure that you are allowed to use the data for your purpose. The redistribution
 56stores the data in the nnU-Net layout, which this module converts to the layout of the official release.
 57Its content was verified against the official release: 50 CT volumes, the label ids 0 - 22, and the same
 58file sizes as the copy of 'Task1_HaN_OAR.zip' that is linked from
 59https://github.com/openmedlab/MedLSAM.
 60
 61NOTE: There is no openly published copy of 'task2', 'task3' and 'task4'. They require registration and
 62cannot be downloaded automatically. Please follow these steps:
 63- Register for the challenge at https://structseg2019.grand-challenge.org and join it, then visit
 64  https://structseg2019.grand-challenge.org/Download/ and follow the instructions of the organizers to
 65  obtain the archives 'Task2_Naso_GTV.zip', 'Task3_Thoracic_OAR.zip' and 'Task4_Lung_GTV.zip' (the GTV
 66  tasks are also deposited at https://doi.org/10.21227/h75x-gt46, which needs a paid IEEE DataPort
 67  subscription and holds only the 83 byte file 'download url.txt' with a download url, not the images).
 68- Extract the archives of the tasks you need into '<path>', such that
 69  '<path>/Naso_GTV/<case_id>/data.nii.gz' and '<path>/Naso_GTV/<case_id>/label.nii.gz' exist
 70  (and equivalently '<path>/Thoracic_OAR/...' and '<path>/Lung_GTV/...').
 71  The case ids are the numbers 1 - 50. An extra enclosing folder, e.g.
 72  '<path>/Task3_Thoracic_OAR/Thoracic_OAR/1/data.nii.gz', is handled as well. A manual download of
 73  'Task1_HaN_OAR.zip' is picked up in the same way and is then used instead of the redistribution.
 74The only openly published parts of these three tasks that could be found are in
 75https://huggingface.co/datasets/blueyo0/SA-Med3D-140K, which is not a usable substitute: it holds the 50
 76task3 volumes with masks for only 4 of the 6 organs (no spinal cord and no trachea), the 50 task2 volumes
 77without any labels, and nothing at all of task4, all resampled to an isotropic spacing of 1.5 mm. Note that
 78its task numbering differs from the challenge: its 'subtask2' entries are task3 and its 'subtask4' entries
 79are task2.
 80
 81The most promising lead for the two GTV tasks is https://rec.ustc.edu.cn/share/b812d430-f577-11ed-a202-03afc6a1d18f,
 82which https://github.com/shijun18/GTV_AutoSeg offers as the code and data of a paper that evaluates on exactly
 83task2 and task4, and whose author also made the IEEE DataPort deposit above, so it is probably the url that
 84the deposit points to. It is not open either: listing that share needs a login at the hosting site.
 85
 86The dataset is located at https://structseg2019.grand-challenge.org.
 87
 88The challenge does not have a dedicated publication. Please cite the challenge website and, for the
 89gross target volume tasks, the data record https://doi.org/10.21227/h75x-gt46 if you use this dataset
 90in your research.
 91"""
 92
 93import os
 94import shutil
 95from glob import glob
 96from natsort import natsorted
 97from typing import Union, Tuple, Literal, List
 98
 99from torch.utils.data import Dataset, DataLoader
100
101import torch_em
102
103from .. import util
104
105
106TASK_FOLDERS = {
107    "task1": "HaN_OAR", "task2": "Naso_GTV", "task3": "Thoracic_OAR", "task4": "Lung_GTV",
108}
109
110# Only 'task1' has an openly published copy, see the module docstring.
111URLS = {"task1": "https://huggingface.co/datasets/Luffy503/VoCo_Downstream/resolve/main/Dataset190_Structseg19.zip"}
112
113CHECKSUMS = {"task1": "0e818cba20bbb783512c6062abc8e2adb43545f5870d630128df1912006c7703"}
114
115MIRROR_FOLDERS = {"task1": "Dataset190_Structseg19"}
116
117HAN_OAR_NAMES = [
118    "brainstem", "eye_l", "eye_r", "lens_l", "lens_r", "optic_nerve_l", "optic_nerve_r", "optic_chiasm",
119    "temporal_lobe_l", "temporal_lobe_r", "pituitary", "parotid_gland_l", "parotid_gland_r", "inner_ear_l",
120    "inner_ear_r", "middle_ear_l", "middle_ear_r", "tm_joint_l", "tm_joint_r", "spinal_cord", "mandible_l",
121    "mandible_r",
122]
123
124THORACIC_OAR_NAMES = ["lung_l", "lung_r", "heart", "esophagus", "trachea", "spinal_cord"]
125
126LABEL_IDS = {
127    "task1": {"background": 0, **{name: i + 1 for i, name in enumerate(HAN_OAR_NAMES)}},
128    "task2": {"background": 0, "nasopharynx_gtv": 1},
129    "task3": {"background": 0, **{name: i + 1 for i, name in enumerate(THORACIC_OAR_NAMES)}},
130    "task4": {"background": 0, "lung_gtv": 1},
131}
132
133
134def _find_task_dir(path, task):
135    folder = TASK_FOLDERS[task]
136    candidates = [os.path.join(path, folder), *glob(os.path.join(path, "*", folder))]
137    for candidate in candidates:
138        if len(glob(os.path.join(candidate, "*", "data.nii.gz"))) > 0:
139            return candidate
140    return None
141
142
143def _convert_mirror_layout(mirror_dir, data_dir):
144    """Convert the nnU-Net layout of the redistribution into the layout of the official release."""
145    image_paths = natsorted(glob(os.path.join(mirror_dir, "imagesTr", "*_0000.nii.gz")))
146    assert len(image_paths) > 0, f"Could not find any images in '{mirror_dir}'."
147
148    for image_path in image_paths:
149        case_id = os.path.basename(image_path)[:-len("_0000.nii.gz")]
150        case_dir = os.path.join(data_dir, case_id)
151        os.makedirs(case_dir, exist_ok=True)
152
153        label_path = os.path.join(mirror_dir, "labelsTr", f"{case_id}.nii.gz")
154        assert os.path.exists(label_path), f"Could not find the labels for case '{case_id}'."
155
156        shutil.move(image_path, os.path.join(case_dir, "data.nii.gz"))
157        shutil.move(label_path, os.path.join(case_dir, "label.nii.gz"))
158
159    shutil.rmtree(mirror_dir)
160
161
162def get_structseg_data(
163    path: Union[os.PathLike, str], task: Literal["task1", "task2", "task3", "task4"], download: bool = False
164) -> str:
165    """Obtain the StructSeg dataset.
166
167    Args:
168        path: Filepath to a folder where the data is stored.
169        task: The task of the challenge. Either 'task1', 'task2', 'task3' or 'task4'.
170        download: Whether to download the data if it is not present.
171
172    Returns:
173        Filepath where the data of the chosen task is stored.
174    """
175    if task not in TASK_FOLDERS:
176        raise ValueError(f"'{task}' is not a valid task. Choose one of {list(TASK_FOLDERS.keys())}.")
177
178    data_dir = _find_task_dir(path, task)
179    if data_dir is not None:
180        return data_dir
181
182    if download and task in URLS:
183        os.makedirs(path, exist_ok=True)
184
185        zip_path = os.path.join(path, os.path.basename(URLS[task]))
186        util.download_source(path=zip_path, url=URLS[task], download=download, checksum=CHECKSUMS[task])
187        util.unzip(zip_path=zip_path, dst=path)
188
189        data_dir = os.path.join(path, TASK_FOLDERS[task])
190        _convert_mirror_layout(os.path.join(path, MIRROR_FOLDERS[task]), data_dir)
191        return data_dir
192
193    msg = f"Could not find the StructSeg '{task}' data ('{TASK_FOLDERS[task]}') at '{path}'. "
194    msg += "'torch_em' cannot download this dataset, as it requires registration for the challenge. "
195    msg += "Please register at https://structseg2019.grand-challenge.org, join the challenge and follow the "
196    msg += "instructions at https://structseg2019.grand-challenge.org/Download/ to obtain the archive "
197    msg += f"'Task{task[-1]}_{TASK_FOLDERS[task]}.zip'. Then extract it into '{path}', such that "
198    msg += f"'{os.path.join(path, TASK_FOLDERS[task], '1', 'data.nii.gz')}' exists."
199    if download:
200        raise NotImplementedError(msg)
201    else:
202        raise FileNotFoundError(msg)
203
204
205def get_structseg_paths(
206    path: Union[os.PathLike, str], task: Literal["task1", "task2", "task3", "task4"], download: bool = False
207) -> Tuple[List[str], List[str]]:
208    """Get paths to the StructSeg data.
209
210    Args:
211        path: Filepath to a folder where the data is stored.
212        task: The task of the challenge. Either 'task1', 'task2', 'task3' or 'task4'.
213        download: Whether to download the data if it is not present.
214
215    Returns:
216        List of filepaths for the image data.
217        List of filepaths for the label data.
218    """
219    data_dir = get_structseg_data(path, task, download)
220
221    raw_paths = natsorted(glob(os.path.join(data_dir, "*", "data.nii.gz")))
222    label_paths = [os.path.join(os.path.dirname(p), "label.nii.gz") for p in raw_paths]
223    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
224
225    return raw_paths, label_paths
226
227
228def get_structseg_dataset(
229    path: Union[os.PathLike, str],
230    patch_shape: Tuple[int, ...],
231    task: Literal["task1", "task2", "task3", "task4"] = "task1",
232    resize_inputs: bool = False,
233    download: bool = False,
234    **kwargs
235) -> Dataset:
236    """Get the StructSeg dataset for organ-at-risk and gross target volume segmentation.
237
238    Args:
239        path: Filepath to a folder where the data is stored.
240        patch_shape: The patch shape to use for training.
241        task: The task of the challenge. Either 'task1', 'task2', 'task3' or 'task4'.
242        resize_inputs: Whether to resize inputs to the desired patch shape.
243        download: Whether to download the data if it is not present.
244        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
245
246    Returns:
247        The segmentation dataset.
248    """
249    raw_paths, label_paths = get_structseg_paths(path, task, download)
250
251    if resize_inputs:
252        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
253        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
254            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
255        )
256
257    return torch_em.default_segmentation_dataset(
258        raw_paths=raw_paths,
259        raw_key="data",
260        label_paths=label_paths,
261        label_key="data",
262        patch_shape=patch_shape,
263        is_seg_dataset=True,
264        **kwargs
265    )
266
267
268def get_structseg_loader(
269    path: Union[os.PathLike, str],
270    batch_size: int,
271    patch_shape: Tuple[int, ...],
272    task: Literal["task1", "task2", "task3", "task4"] = "task1",
273    resize_inputs: bool = False,
274    download: bool = False,
275    **kwargs
276) -> DataLoader:
277    """Get the StructSeg dataloader for organ-at-risk and gross target volume segmentation.
278
279    Args:
280        path: Filepath to a folder where the data is stored.
281        batch_size: The batch size for training.
282        patch_shape: The patch shape to use for training.
283        task: The task of the challenge. Either 'task1', 'task2', 'task3' or 'task4'.
284        resize_inputs: Whether to resize inputs to the desired patch shape.
285        download: Whether to download the data if it is not present.
286        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
287
288    Returns:
289        The DataLoader.
290    """
291    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
292    dataset = get_structseg_dataset(path, patch_shape, task, resize_inputs, download, **ds_kwargs)
293    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)