torch_em.data.datasets.medical.prostate158

  1"""The Prostate158 dataset contains annotations for prostate zone and prostate cancer segmentation
  2in biparametric 3T MRI.
  3
  4Each study provides a T2-weighted sequence ('t2'), a diffusion-weighted sequence ('dwi') and the corresponding
  5apparent diffusion coefficient map ('adc'). All sequences of a study were resampled to the same grid,
  6so that the labels can be used with any of them.
  7
  8NOTE: The label legends are described as following:
  91: For the anatomical zones ('anatomy', annotated by reader 1 in the T2-weighted sequence):
 10- background: 0, transition zone (central gland): 1 and peripheral zone: 2.
 112: For the prostate cancer lesions ('tumor', annotated by reader 1 in the ADC map):
 12- background: 0 and tumor: 1. Studies without a lesion have an empty label volume.
 13
 14The official split provides 119 training and 20 validation studies (record https://zenodo.org/records/6481141)
 15and 19 test studies with additional annotations from a second reader (record https://zenodo.org/records/6592345).
 16
 17The dataset is located at https://github.com/kbressem/prostate158.
 18
 19This dataset is from the publication https://doi.org/10.1016/j.compbiomed.2022.105817.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24import csv
 25from typing import Union, Tuple, Literal, List
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URLS = {
 35    "train": "https://zenodo.org/records/6481141/files/prostate158_train.zip?download=1",
 36    "test": "https://zenodo.org/records/6592345/files/prostate158_test.zip?download=1",
 37}
 38
 39CHECKSUMS = {
 40    "train": "7a97b263be1bdbc79f6c8a3461e010b2cea9a266249af47671616dadff77da2e",
 41    "test": "d64a7b94f0654de21150af9b7a5734b2921ec4793bc129b98803ef9ae9d459c4",
 42}
 43
 44LABEL_COLUMNS = {"anatomy": "t2_anatomy_reader1", "tumor": "adc_tumor_reader1"}
 45
 46
 47def get_prostate158_data(
 48    path: Union[os.PathLike, str], split: Literal["train", "valid", "test"], download: bool = False
 49) -> str:
 50    """Download the Prostate158 dataset.
 51
 52    Args:
 53        path: Filepath to a folder where the data is downloaded for further processing.
 54        split: The choice of data split. The 'train' and 'valid' splits share one archive, 'test' has its own.
 55        download: Whether to download the data if it is not present.
 56
 57    Returns:
 58        Filepath to the folder with the data of the requested split.
 59    """
 60    if split not in ("train", "valid", "test"):
 61        raise ValueError(f"'{split}' is not a valid split.")
 62
 63    archive = "test" if split == "test" else "train"
 64    data_dir = os.path.join(path, f"prostate158_{archive}")
 65    if os.path.exists(data_dir):
 66        return data_dir
 67
 68    os.makedirs(path, exist_ok=True)
 69
 70    zip_path = os.path.join(path, f"prostate158_{archive}.zip")
 71    util.download_source(path=zip_path, url=URLS[archive], download=download, checksum=CHECKSUMS[archive])
 72    util.unzip(zip_path=zip_path, dst=path)
 73
 74    return data_dir
 75
 76
 77def get_prostate158_paths(
 78    path: Union[os.PathLike, str],
 79    split: Literal["train", "valid", "test"],
 80    sequence: Literal["t2", "adc", "dwi"] = "t2",
 81    label_type: Literal["anatomy", "tumor"] = "anatomy",
 82    download: bool = False,
 83) -> Tuple[List[str], List[str]]:
 84    """Get paths to the Prostate158 data.
 85
 86    Args:
 87        path: Filepath to a folder where the data is downloaded for further processing.
 88        split: The choice of data split. Either 'train', 'valid' or 'test'.
 89        sequence: The MRI sequence to use as input. Either 't2', 'adc' or 'dwi'.
 90        label_type: The type of annotations. Either 'anatomy' (prostate zones) or 'tumor' (cancer lesions).
 91        download: Whether to download the data if it is not present.
 92
 93    Returns:
 94        List of filepaths for the image data.
 95        List of filepaths for the label data.
 96    """
 97    if sequence not in ("t2", "adc", "dwi"):
 98        raise ValueError(f"'{sequence}' is not a valid sequence.")
 99    if label_type not in LABEL_COLUMNS:
100        raise ValueError(f"'{label_type}' is not a valid label type.")
101
102    data_dir = get_prostate158_data(path, split, download)
103
104    # The official csv files list the volumes per study, including empty tumor labels for studies without lesion.
105    with open(os.path.join(data_dir, f"{split}.csv"), "r") as f:
106        rows = list(csv.DictReader(f))
107
108    raw_paths = [os.path.join(data_dir, row[sequence]) for row in rows]
109    label_paths = [os.path.join(data_dir, row[LABEL_COLUMNS[label_type]]) for row in rows]
110
111    missing = [p for p in raw_paths + label_paths if not os.path.exists(p)]
112    assert len(missing) == 0, f"The following files are missing: {missing}"
113
114    return raw_paths, label_paths
115
116
117def get_prostate158_dataset(
118    path: Union[os.PathLike, str],
119    patch_shape: Tuple[int, ...],
120    split: Literal["train", "valid", "test"],
121    sequence: Literal["t2", "adc", "dwi"] = "t2",
122    label_type: Literal["anatomy", "tumor"] = "anatomy",
123    resize_inputs: bool = False,
124    download: bool = False,
125    **kwargs
126) -> Dataset:
127    """Get the Prostate158 dataset for prostate zone and prostate cancer segmentation.
128
129    Args:
130        path: Filepath to a folder where the data is downloaded for further processing.
131        patch_shape: The patch shape to use for training.
132        split: The choice of data split. Either 'train', 'valid' or 'test'.
133        sequence: The MRI sequence to use as input. Either 't2', 'adc' or 'dwi'.
134        label_type: The type of annotations. Either 'anatomy' (prostate zones) or 'tumor' (cancer lesions).
135        resize_inputs: Whether to resize inputs to the desired patch shape.
136        download: Whether to download the data if it is not present.
137        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
138
139    Returns:
140        The segmentation dataset.
141    """
142    raw_paths, label_paths = get_prostate158_paths(path, split, sequence, label_type, download)
143
144    if resize_inputs:
145        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
146        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
147            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
148        )
149
150    return torch_em.default_segmentation_dataset(
151        raw_paths=raw_paths,
152        raw_key="data",
153        label_paths=label_paths,
154        label_key="data",
155        patch_shape=patch_shape,
156        is_seg_dataset=True,
157        **kwargs
158    )
159
160
161def get_prostate158_loader(
162    path: Union[os.PathLike, str],
163    batch_size: int,
164    patch_shape: Tuple[int, ...],
165    split: Literal["train", "valid", "test"],
166    sequence: Literal["t2", "adc", "dwi"] = "t2",
167    label_type: Literal["anatomy", "tumor"] = "anatomy",
168    resize_inputs: bool = False,
169    download: bool = False,
170    **kwargs
171) -> DataLoader:
172    """Get the Prostate158 dataloader for prostate zone and prostate cancer segmentation.
173
174    Args:
175        path: Filepath to a folder where the data is downloaded for further processing.
176        batch_size: The batch size for training.
177        patch_shape: The patch shape to use for training.
178        split: The choice of data split. Either 'train', 'valid' or 'test'.
179        sequence: The MRI sequence to use as input. Either 't2', 'adc' or 'dwi'.
180        label_type: The type of annotations. Either 'anatomy' (prostate zones) or 'tumor' (cancer lesions).
181        resize_inputs: Whether to resize inputs to the desired patch shape.
182        download: Whether to download the data if it is not present.
183        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
184
185    Returns:
186        The DataLoader.
187    """
188    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
189    dataset = get_prostate158_dataset(
190        path, patch_shape, split, sequence, label_type, resize_inputs, download, **ds_kwargs
191    )
192    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)