torch_em.data.datasets.medical.openkbp

  1"""The OpenKBP dataset contains annotations for organ-at-risk (OAR) and planning target volume (PTV)
  2segmentation in radiotherapy planning CT scans of head-and-neck cancer patients.
  3
  4The data was curated for the OpenKBP Grand Challenge (https://github.com/ababier/open-kbp), which was held
  5at AAPM 2020 and targets models for knowledge-based dose prediction. The public release consists of 340
  6patients treated with intensity modulated radiation therapy, split into the official 200 training, 40
  7validation and 100 test studies (selected with the 'split' argument). Each patient has a 128x128x128 CT
  8volume together with up to 10 structure masks: 7 OARs ('Brainstem', 'SpinalCord', 'RightParotid',
  9'LeftParotid', 'Esophagus', 'Larynx', 'Mandible') and 3 PTVs ('PTV56', 'PTV63', 'PTV70'), see `LABEL_IDS`.
 10Not every patient has every structure delineated; missing structures are simply not present in that
 11patient's label volume.
 12
 13NOTE: The raw data also provides a per-voxel dose distribution and a feasible dose mask for the OpenKBP
 14dose prediction challenge. This module only exposes the OAR/PTV structure masks as segmentation targets;
 15dose regression is out of scope for this library.
 16
 17The data is hosted directly (as plain csv files, not through git-lfs) in the challenge repository at
 18https://github.com/ababier/open-kbp. This module downloads it from a pinned commit of that repository, so
 19that the checksum of the downloaded archive stays reproducible. It is distributed for research use; see the
 20repository for the exact license terms.
 21
 22The CT and structure files each store a flattened (raveled, C order) sparse representation of the
 23128x128x128 patient volume: only the non-background voxel indices (and, for the CT scan, their values) are
 24listed. This module reconstructs the dense CT volume and a single-channel structure label volume (using the
 25priority order of `LABEL_IDS`, so that a PTV overwrites an OAR at an overlapping voxel) and stores them as
 26hdf5 volumes (keys 'raw' and 'labels') for efficient access.
 27
 28This dataset is from the publication https://doi.org/10.1002/mp.14845.
 29Please cite it if you use this dataset in your research.
 30"""
 31
 32import os
 33from glob import glob
 34from tqdm import tqdm
 35from natsort import natsorted
 36from typing import Union, Tuple, List, Literal
 37
 38import numpy as np
 39
 40from torch.utils.data import Dataset, DataLoader
 41
 42import torch_em
 43
 44from .. import util
 45
 46
 47URL = "https://codeload.github.com/ababier/open-kbp/zip/ce625e62f3b04203f22bd9d1634f3e8fb0245e46"
 48CHECKSUM = "4074858f3f0d9d7f639342349e2799c1109f3018e3bd487827f334b591ebf697"
 49
 50COMMIT = "ce625e62f3b04203f22bd9d1634f3e8fb0245e46"
 51
 52SPLITS = {"train": "train-pats", "val": "validation-pats", "test": "test-pats"}
 53
 54N_PATIENTS = {"train": 200, "val": 40, "test": 100}
 55
 56PATIENT_SHAPE = (128, 128, 128)
 57
 58OARS = ["Brainstem", "SpinalCord", "RightParotid", "LeftParotid", "Esophagus", "Larynx", "Mandible"]
 59TARGETS = ["PTV56", "PTV63", "PTV70"]
 60STRUCTURES = OARS + TARGETS
 61
 62LABEL_IDS = {"background": 0}
 63LABEL_IDS.update({name: i + 1 for i, name in enumerate(STRUCTURES)})
 64
 65
 66def _load_sparse_csv(csv_path):
 67    import pandas as pd
 68
 69    df = pd.read_csv(csv_path, index_col=0)
 70    indices = df.index.to_numpy()
 71    if df["data"].isna().all():  # This is a binary structure mask, so it only stores the indices.
 72        return indices, None
 73    else:  # This is a sparse volume (eg. the CT scan), so it also stores a value per index.
 74        return indices, df["data"].to_numpy()
 75
 76
 77def _preprocess_patient(patient_dir, volume_path):
 78    import h5py
 79
 80    raw = np.zeros(np.prod(PATIENT_SHAPE), dtype="int16")
 81    indices, values = _load_sparse_csv(os.path.join(patient_dir, "ct.csv"))
 82    raw[indices] = values
 83    raw = raw.reshape(PATIENT_SHAPE)
 84
 85    labels = np.zeros(np.prod(PATIENT_SHAPE), dtype="uint8")
 86    for structure in STRUCTURES:
 87        structure_path = os.path.join(patient_dir, f"{structure}.csv")
 88        if not os.path.exists(structure_path):
 89            continue
 90        indices, _ = _load_sparse_csv(structure_path)
 91        labels[indices] = LABEL_IDS[structure]
 92    labels = labels.reshape(PATIENT_SHAPE)
 93
 94    with h5py.File(f"{volume_path}.tmp", "w") as f:
 95        f.create_dataset("raw", data=raw, compression="gzip")
 96        f.create_dataset("labels", data=labels, compression="gzip")
 97
 98    os.rename(f"{volume_path}.tmp", volume_path)
 99
100
101def _preprocess_inputs(data_dir, split, preprocessed_dir):
102    patient_dirs = natsorted(glob(os.path.join(data_dir, SPLITS[split], "pt_*")))
103    os.makedirs(preprocessed_dir, exist_ok=True)
104    for patient_dir in tqdm(patient_dirs, desc=f"Preprocessing the OpenKBP '{split}' patients"):
105        volume_path = os.path.join(preprocessed_dir, f"{os.path.basename(patient_dir)}.h5")
106        if os.path.exists(volume_path):
107            continue
108        _preprocess_patient(patient_dir, volume_path)
109
110
111def get_openkbp_data(path: Union[os.PathLike, str], download: bool = False) -> str:
112    """Download the OpenKBP dataset.
113
114    Args:
115        path: Filepath to a folder where the data is downloaded for further processing.
116        download: Whether to download the data if it is not present.
117
118    Returns:
119        Filepath where the provided data is stored.
120    """
121    data_dir = os.path.join(path, f"open-kbp-{COMMIT}", "provided-data")
122    if os.path.exists(data_dir):
123        return data_dir
124
125    os.makedirs(path, exist_ok=True)
126
127    zip_path = os.path.join(path, "open-kbp.zip")
128    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
129    util.unzip(zip_path=zip_path, dst=path)
130
131    return data_dir
132
133
134def get_openkbp_paths(
135    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
136) -> List[str]:
137    """Get paths to the OpenKBP data.
138
139    Args:
140        path: Filepath to a folder where the data is downloaded for further processing.
141        split: The choice of data split. Either 'train', 'val' or 'test'.
142        download: Whether to download the data if it is not present.
143
144    Returns:
145        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
146    """
147    if split not in SPLITS:
148        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(SPLITS.keys())}.")
149
150    preprocessed_dir = os.path.join(path, "preprocessed", split)
151    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) != N_PATIENTS[split]:
152        data_dir = get_openkbp_data(path, download)
153        _preprocess_inputs(data_dir, split, preprocessed_dir)
154
155    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
156    assert len(volume_paths) == N_PATIENTS[split]
157
158    return volume_paths
159
160
161def get_openkbp_dataset(
162    path: Union[os.PathLike, str],
163    patch_shape: Tuple[int, ...],
164    split: Literal["train", "val", "test"],
165    resize_inputs: bool = False,
166    download: bool = False,
167    **kwargs
168) -> Dataset:
169    """Get the OpenKBP dataset for organ-at-risk and planning target volume segmentation.
170
171    Args:
172        path: Filepath to a folder where the data is downloaded for further processing.
173        patch_shape: The patch shape to use for training.
174        split: The choice of data split. Either 'train', 'val' or 'test'.
175        resize_inputs: Whether to resize inputs to the desired patch shape.
176        download: Whether to download the data if it is not present.
177        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
178
179    Returns:
180        The segmentation dataset.
181    """
182    volume_paths = get_openkbp_paths(path, split, download)
183
184    if resize_inputs:
185        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
186        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
187            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
188        )
189
190    return torch_em.default_segmentation_dataset(
191        raw_paths=volume_paths,
192        raw_key="raw",
193        label_paths=volume_paths,
194        label_key="labels",
195        patch_shape=patch_shape,
196        is_seg_dataset=True,
197        **kwargs
198    )
199
200
201def get_openkbp_loader(
202    path: Union[os.PathLike, str],
203    batch_size: int,
204    patch_shape: Tuple[int, ...],
205    split: Literal["train", "val", "test"],
206    resize_inputs: bool = False,
207    download: bool = False,
208    **kwargs
209) -> DataLoader:
210    """Get the OpenKBP dataloader for organ-at-risk and planning target volume segmentation.
211
212    Args:
213        path: Filepath to a folder where the data is downloaded for further processing.
214        batch_size: The batch size for training.
215        patch_shape: The patch shape to use for training.
216        split: The choice of data split. Either 'train', 'val' or 'test'.
217        resize_inputs: Whether to resize inputs to the desired patch shape.
218        download: Whether to download the data if it is not present.
219        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
220
221    Returns:
222        The DataLoader.
223    """
224    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
225    dataset = get_openkbp_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
226    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)