torch_em.data.datasets.medical.pddca

  1"""The PDDCA dataset contains annotations for organs-at-risk segmentation in head and neck CT.
  2
  3The dataset (Public Domain Database for Computational Anatomy, version 1.4.1) consists of 48 CT scans from the
  4RTOG 0522 clinical trial with manual segmentations of nine organs at risk. It was used for the
  5MICCAI 2015 Head and Neck Auto-Segmentation Challenge.
  6
  7NOTE: The per-structure binary masks are combined into one semantic label volume with the following ids:
  8- background: 0
  9- brainstem: 1
 10- optic chiasm: 2
 11- mandible: 3
 12- left optic nerve: 4
 13- right optic nerve: 5
 14- left parotid gland: 6
 15- right parotid gland: 7
 16- left submandibular gland: 8
 17- right submandibular gland: 9
 18A few structures overlap by a handful of voxels. In this case, the smaller structure takes priority:
 19the masks are written in the order mandible, parotids, submandibular glands, brainstem, optic nerves, chiasm,
 20so that a later structure overwrites an earlier one. Not all structures are annotated for every patient
 21(e.g. some patients lack the mandible or the right submandibular gland); missing structures are left as background.
 22
 23The official challenge sub-packages are exposed as splits: 'train' (25 scans, 0522c0001 to 0522c0328),
 24'train_additional' (8 scans, 0522c0329 to 0522c0479), 'test_offsite' (10 scans, 0522c0555 to 0522c0746)
 25and 'test_onsite' (5 scans, 0522c0788 to 0522c0878).
 26
 27The dataset is located at https://www.imagenglab.com/newsite/pddca/.
 28
 29This dataset is from the publication https://doi.org/10.1002/mp.12197.
 30Please cite it if you use this dataset in your research.
 31"""
 32
 33import os
 34from glob import glob
 35from tqdm import tqdm
 36from natsort import natsorted
 37from typing import Union, Tuple, Literal, List, Optional
 38
 39import numpy as np
 40
 41from torch.utils.data import Dataset, DataLoader
 42
 43import torch_em
 44
 45from .. import util
 46
 47
 48URLS = [
 49    "https://www.imagenglab.com/data/pddca/PDDCA-1.4.1_part1.zip",
 50    "https://www.imagenglab.com/data/pddca/PDDCA-1.4.1_part2.zip",
 51    "https://www.imagenglab.com/data/pddca/PDDCA-1.4.1_part3.zip",
 52]
 53
 54CHECKSUMS = [
 55    "5b47b94f1e6aaad0a10c694906bd40ee8bcfc701e6d3bcbe3f12005815b721ef",
 56    "a5354e5373050f958e2696088a24f25f72361844d8b0d6988c16a3c989fd5659",
 57    "58bc70db301d3f0fd2e856a33e7d2be929859bc5b1a37d7c83a47f2bbe6a7fb6",
 58]
 59
 60LABEL_IDS = {
 61    "BrainStem": 1,
 62    "Chiasm": 2,
 63    "Mandible": 3,
 64    "OpticNerve_L": 4,
 65    "OpticNerve_R": 5,
 66    "Parotid_L": 6,
 67    "Parotid_R": 7,
 68    "Submandibular_L": 8,
 69    "Submandibular_R": 9,
 70}
 71
 72# The order in which the structures are written to the label volume. Later structures overwrite earlier ones.
 73WRITE_ORDER = [
 74    "Mandible", "Parotid_L", "Parotid_R", "Submandibular_L", "Submandibular_R",
 75    "BrainStem", "OpticNerve_L", "OpticNerve_R", "Chiasm",
 76]
 77
 78SPLITS = {
 79    "train": (1, 328),
 80    "train_additional": (329, 479),
 81    "test_offsite": (555, 746),
 82    "test_onsite": (788, 878),
 83}
 84
 85
 86def _convert_case(case_dir, out_path):
 87    import nrrd
 88    import h5py
 89
 90    raw, header = nrrd.read(os.path.join(case_dir, "img.nrrd"))
 91    labels = np.zeros(raw.shape, dtype="uint8")
 92    for name in WRITE_ORDER:
 93        mask_path = os.path.join(case_dir, "structures", f"{name}.nrrd")
 94        if not os.path.exists(mask_path):
 95            continue
 96        mask, _ = nrrd.read(mask_path)
 97        assert mask.shape == raw.shape, f"Shape mismatch for {mask_path}."
 98        labels[mask > 0] = LABEL_IDS[name]
 99
100    # The nrrd arrays are stored in (x, y, z) order, we transpose them to (z, y, x).
101    raw, labels = raw.transpose(2, 1, 0), labels.transpose(2, 1, 0)
102    with h5py.File(out_path, "w") as f:
103        f.create_dataset("raw", data=raw, compression="gzip")
104        f.create_dataset("labels", data=labels, compression="gzip")
105        f.attrs["spacing"] = np.diag(header["space directions"])[::-1].tolist()
106
107
108def get_pddca_data(path: Union[os.PathLike, str], download: bool = False) -> str:
109    """Download the PDDCA dataset and convert it to hdf5 volumes.
110
111    Args:
112        path: Filepath to a folder where the data is downloaded for further processing.
113        download: Whether to download the data if it is not present.
114
115    Returns:
116        Filepath to the folder with the preprocessed hdf5 volumes.
117    """
118    data_dir = os.path.join(path, "preprocessed")
119    if os.path.exists(data_dir):
120        return data_dir
121
122    os.makedirs(path, exist_ok=True)
123    raw_dir = os.path.join(path, "raw")
124    for url, checksum in zip(URLS, CHECKSUMS):
125        zip_path = os.path.join(path, os.path.basename(url))
126        util.download_source(path=zip_path, url=url, download=download, checksum=checksum)
127        util.unzip(zip_path=zip_path, dst=raw_dir)
128
129    case_dirs = natsorted(glob(os.path.join(raw_dir, "0522c*")))
130    assert len(case_dirs) == 48, f"Expected 48 cases, found {len(case_dirs)}."
131
132    os.makedirs(data_dir, exist_ok=True)
133    for case_dir in tqdm(case_dirs, desc="Converting PDDCA volumes to hdf5"):
134        _convert_case(case_dir, os.path.join(data_dir, f"{os.path.basename(case_dir)}.h5"))
135
136    return data_dir
137
138
139def get_pddca_paths(
140    path: Union[os.PathLike, str],
141    split: Optional[Literal["train", "train_additional", "test_offsite", "test_onsite"]] = None,
142    download: bool = False,
143) -> List[str]:
144    """Get paths to the PDDCA data.
145
146    Args:
147        path: Filepath to a folder where the data is downloaded for further processing.
148        split: The choice of data split. One of 'train', 'train_additional', 'test_offsite' or 'test_onsite'.
149            By default, all volumes are returned.
150        download: Whether to download the data if it is not present.
151
152    Returns:
153        List of filepaths for the hdf5 volumes with image and label data.
154    """
155    data_dir = get_pddca_data(path, download)
156    volume_paths = natsorted(glob(os.path.join(data_dir, "0522c*.h5")))
157
158    if split is not None:
159        if split not in SPLITS:
160            raise ValueError(f"'{split}' is not a valid split. Choose one of {list(SPLITS)}.")
161        lower, upper = SPLITS[split]
162        case_ids = [int(os.path.basename(p).replace("0522c", "").replace(".h5", "")) for p in volume_paths]
163        volume_paths = [p for p, case_id in zip(volume_paths, case_ids) if lower <= case_id <= upper]
164
165    assert len(volume_paths) > 0
166    return volume_paths
167
168
169def get_pddca_dataset(
170    path: Union[os.PathLike, str],
171    patch_shape: Tuple[int, ...],
172    split: Optional[Literal["train", "train_additional", "test_offsite", "test_onsite"]] = None,
173    resize_inputs: bool = False,
174    download: bool = False,
175    **kwargs
176) -> Dataset:
177    """Get the PDDCA dataset for organs-at-risk segmentation in head and neck CT.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        patch_shape: The patch shape to use for training.
182        split: The choice of data split. One of 'train', 'train_additional', 'test_offsite' or 'test_onsite'.
183            By default, all volumes are returned.
184        resize_inputs: Whether to resize inputs to the desired patch shape.
185        download: Whether to download the data if it is not present.
186        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
187
188    Returns:
189        The segmentation dataset.
190    """
191    volume_paths = get_pddca_paths(path, split, download)
192
193    if resize_inputs:
194        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
195        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
196            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
197        )
198
199    return torch_em.default_segmentation_dataset(
200        raw_paths=volume_paths,
201        raw_key="raw",
202        label_paths=volume_paths,
203        label_key="labels",
204        patch_shape=patch_shape,
205        is_seg_dataset=True,
206        **kwargs
207    )
208
209
210def get_pddca_loader(
211    path: Union[os.PathLike, str],
212    batch_size: int,
213    patch_shape: Tuple[int, ...],
214    split: Optional[Literal["train", "train_additional", "test_offsite", "test_onsite"]] = None,
215    resize_inputs: bool = False,
216    download: bool = False,
217    **kwargs
218) -> DataLoader:
219    """Get the PDDCA dataloader for organs-at-risk segmentation in head and neck CT.
220
221    Args:
222        path: Filepath to a folder where the data is downloaded for further processing.
223        batch_size: The batch size for training.
224        patch_shape: The patch shape to use for training.
225        split: The choice of data split. One of 'train', 'train_additional', 'test_offsite' or 'test_onsite'.
226            By default, all volumes are returned.
227        resize_inputs: Whether to resize inputs to the desired patch shape.
228        download: Whether to download the data if it is not present.
229        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
230
231    Returns:
232        The DataLoader.
233    """
234    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
235    dataset = get_pddca_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
236    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)