torch_em.data.datasets.medical.nci_isbi_prostate

  1"""The NCI-ISBI 2013 Prostate dataset contains annotations for prostate zone segmentation in T2-weighted MRI.
  2
  3The data comes from the 'NCI-ISBI 2013 Challenge: Automated Segmentation of Prostate Structures' and consists of
  480 T2-weighted MRI of the prostate, 40 acquired at 3T with a surface coil (the TCIA collection 'Prostate-3T') and
  540 acquired at 1.5T with an endorectal coil (the TCIA collection 'PROSTATE-DIAGNOSIS'). The official challenge split
  6is exposed via the 'split' argument: 'train' (60 volumes), 'leaderboard' (10 volumes) and 'test' (10 volumes).
  7
  8The MRI are distributed as DICOM series on TCIA and the manual segmentations as NRRD files, which are converted and
  9stored together in hdf5 files by this module. The label ids are: 0 = background, 1 = peripheral zone,
 102 = central gland.
 11
 12NOTE: This requires the pydicom and pynrrd python packages.
 13
 14The dataset is located at https://www.cancerimagingarchive.net/analysis-result/isbi-mr-prostate-2013/.
 15
 16This dataset is from the publication https://doi.org/10.7937/K9/TCIA.2015.zF0vlOPv.
 17Please cite it if you use this dataset in your research.
 18"""
 19
 20import os
 21import csv
 22from glob import glob
 23from tqdm import tqdm
 24from natsort import natsorted
 25from typing import Union, Tuple, List, Literal
 26
 27import numpy as np
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36BASE_URL = "https://www.cancerimagingarchive.net/wp-content/uploads/"
 37
 38URLS = {
 39    "train": {
 40        "images": BASE_URL + "ISBI-Prostate-Challenge-Training.tcia",
 41        "labels": BASE_URL + "NCI-ISBI-2013-Prostate-Challenge-Training.zip",
 42    },
 43    "leaderboard": {
 44        "images": BASE_URL + "ISBI-Prostate-Challenge-LeaderBoard.tcia",
 45        "labels": BASE_URL + "NCI-ISBI-2013-Prostate-Challenge-Leaderboard.zip",
 46    },
 47    "test": {
 48        "images": BASE_URL + "ISBI-Prostate-Challenge-Testing.tcia",
 49        "labels": BASE_URL + "NCI-ISBI-2013-Prostate-Challenge-Test.zip",
 50    },
 51}
 52
 53# The DICOM series are downloaded individually from TCIA, so only the label archives have a checksum.
 54CHECKSUMS = {
 55    "train": "c3436559b474c60e78633ea98601241f39cb27e30fb9d48b1d29f2821bfbf047",
 56    "leaderboard": "011945fdd273b2f43f6ba2de8904ad8625800b439dff9d228810b645c65fa0be",
 57    "test": "53c2969c035974de47bcb2d3fa20c1342c99096faf2fc87a0097ed62e89d32e3",
 58}
 59
 60LABEL_IDS = {"background": 0, "peripheral_zone": 1, "central_gland": 2}
 61
 62SPLITS = {"train": "Training", "leaderboard": "Leaderboard", "test": "Test"}
 63
 64N_VOLUMES = {"train": 60, "leaderboard": 10, "test": 10}
 65
 66
 67def _get_dicom_affine(geometry):
 68    """Build the affine that maps the voxel indices (z, y, x) of a DICOM volume to patient coordinates."""
 69    origin = geometry["origin"]
 70    if len(origin) > 1:
 71        slice_step = (origin[-1] - origin[0]) / (len(origin) - 1)
 72    else:
 73        slice_step = np.cross(geometry["row_direction"], geometry["column_direction"])
 74
 75    affine = np.eye(4)
 76    affine[:3, 0] = slice_step
 77    affine[:3, 1] = geometry["column_direction"] * geometry["spacing"][0]
 78    affine[:3, 2] = geometry["row_direction"] * geometry["spacing"][1]
 79    affine[:3, 3] = origin[0]
 80    return affine
 81
 82
 83def _load_nrrd_on_dicom_grid(nrrd_path, geometry, shape):
 84    """Load a NRRD label volume and resample it onto the voxel grid of the DICOM series.
 85
 86    Most segmentations are stored on exactly the grid of the DICOM series, but with a different axis order,
 87    while some (mostly in the 'PROSTATE-DIAGNOSIS' part) are stored on a grid with a different slice spacing.
 88    Both cases are handled by mapping each voxel of the DICOM grid to patient coordinates via the DICOM affine
 89    and back to a NRRD voxel index via the inverse of the NRRD affine, i.e. by nearest neighbor resampling.
 90    """
 91    import nrrd
 92
 93    labels, header = nrrd.read(nrrd_path)
 94    directions = np.asarray(header["space directions"], dtype="float64")
 95    origin = np.asarray(header["space origin"], dtype="float64")
 96    if header.get("space", "").lower().startswith("right-anterior"):  # Convert RAS to the LPS used by DICOM.
 97        directions = directions * np.array([-1.0, -1.0, 1.0])
 98        origin = origin * np.array([-1.0, -1.0, 1.0])
 99
100    nrrd_affine = np.eye(4)
101    nrrd_affine[:3, :3] = directions.T  # The rows of 'space directions' are the directions of the NRRD axes.
102    nrrd_affine[:3, 3] = origin
103
104    to_nrrd_index = np.linalg.inv(nrrd_affine) @ _get_dicom_affine(geometry)
105    resampled = np.zeros(shape, dtype="uint8")
106    yy, xx = np.meshgrid(np.arange(shape[1]), np.arange(shape[2]), indexing="ij")
107    for z in range(shape[0]):
108        target_indices = np.stack([np.full(yy.size, z), yy.ravel(), xx.ravel(), np.ones(yy.size)])
109        indices = np.round(to_nrrd_index[:3] @ target_indices).astype("int")
110        valid = np.all((indices >= 0) & (indices < np.array(labels.shape)[:, None]), axis=0)
111        resampled[z].flat[valid] = labels[tuple(indices[:, valid])]
112
113    return resampled
114
115
116def _preprocess_inputs(dicom_dir, csv_path, label_dir, preprocessed_dir):
117    import h5py
118
119    with open(csv_path, "r") as f:
120        series_uids = {row["Subject ID"]: row["Series UID"] for row in csv.DictReader(f)}
121
122    label_paths = {}
123    for label_path in natsorted(glob(os.path.join(label_dir, "*", "*.nrrd"))):
124        subject_id = os.path.basename(label_path).split(".")[0]
125        # Some label files carry a suffix, e.g. 'ProstateDx-01-0006_correctedLabels.nrrd'.
126        label_paths[subject_id.split("_")[0]] = label_path
127
128    os.makedirs(preprocessed_dir, exist_ok=True)
129    for subject_id, label_path in tqdm(sorted(label_paths.items()), desc="Preprocess NCI-ISBI 2013 Prostate"):
130        volume_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
131        if os.path.exists(volume_path):
132            continue
133
134        raw, geometry = util.load_dicom_series(os.path.join(dicom_dir, series_uids[subject_id]))
135        raw = np.round(raw).astype("int16")
136        labels = _load_nrrd_on_dicom_grid(label_path, geometry, raw.shape)
137
138        with h5py.File(f"{volume_path}.tmp", "w") as f:
139            f.create_dataset("raw", data=raw, compression="gzip")
140            f.create_dataset("labels", data=labels, compression="gzip")
141
142        os.rename(f"{volume_path}.tmp", volume_path)
143
144
145def get_nci_isbi_prostate_data(
146    path: Union[os.PathLike, str], split: Literal["train", "leaderboard", "test"], download: bool = False
147) -> str:
148    """Download the NCI-ISBI 2013 Prostate dataset.
149
150    Args:
151        path: Filepath to a folder where the data is downloaded for further processing.
152        split: The choice of data split. Either 'train', 'leaderboard' or 'test'.
153        download: Whether to download the data if it is not present.
154
155    Returns:
156        Filepath where the preprocessed data is stored.
157    """
158    if split not in SPLITS:
159        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(SPLITS.keys())}.")
160
161    preprocessed_dir = os.path.join(path, "preprocessed", split)
162    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_VOLUMES[split]:
163        return preprocessed_dir
164
165    os.makedirs(path, exist_ok=True)
166
167    label_dir = os.path.join(path, "labels", split)
168    if not os.path.exists(label_dir):
169        zip_path = os.path.join(path, f"{SPLITS[split]}.zip")
170        util.download_source(path=zip_path, url=URLS[split]["labels"], download=download, checksum=CHECKSUMS[split])
171        util.unzip(zip_path=zip_path, dst=label_dir)
172
173    dicom_dir = os.path.join(path, "dicom", split)
174    csv_filename = os.path.join(path, f"{split}_metadata")
175    if not os.path.exists(f"{csv_filename}.csv"):
176        util.download_source_tcia(
177            path=os.path.join(path, f"{split}.tcia"),
178            url=URLS[split]["images"],
179            dst=dicom_dir,
180            csv_filename=csv_filename,
181            download=download,
182        )
183
184    _preprocess_inputs(dicom_dir, f"{csv_filename}.csv", label_dir, preprocessed_dir)
185    return preprocessed_dir
186
187
188def get_nci_isbi_prostate_paths(
189    path: Union[os.PathLike, str], split: Literal["train", "leaderboard", "test"], download: bool = False
190) -> List[str]:
191    """Get paths to the NCI-ISBI 2013 Prostate data.
192
193    Args:
194        path: Filepath to a folder where the data is downloaded for further processing.
195        split: The choice of data split. Either 'train', 'leaderboard' or 'test'.
196        download: Whether to download the data if it is not present.
197
198    Returns:
199        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
200    """
201    data_dir = get_nci_isbi_prostate_data(path, split, download)
202    return natsorted(glob(os.path.join(data_dir, "*.h5")))
203
204
205def get_nci_isbi_prostate_dataset(
206    path: Union[os.PathLike, str],
207    patch_shape: Tuple[int, ...],
208    split: Literal["train", "leaderboard", "test"],
209    resize_inputs: bool = False,
210    download: bool = False,
211    **kwargs
212) -> Dataset:
213    """Get the NCI-ISBI 2013 Prostate dataset for prostate zone segmentation.
214
215    Args:
216        path: Filepath to a folder where the data is downloaded for further processing.
217        patch_shape: The patch shape to use for training.
218        split: The choice of data split. Either 'train', 'leaderboard' or 'test'.
219        resize_inputs: Whether to resize inputs to the desired patch shape.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
222
223    Returns:
224        The segmentation dataset.
225    """
226    volume_paths = get_nci_isbi_prostate_paths(path, split, download)
227
228    if resize_inputs:
229        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
230        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
231            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
232        )
233
234    return torch_em.default_segmentation_dataset(
235        raw_paths=volume_paths,
236        raw_key="raw",
237        label_paths=volume_paths,
238        label_key="labels",
239        patch_shape=patch_shape,
240        is_seg_dataset=True,
241        **kwargs
242    )
243
244
245def get_nci_isbi_prostate_loader(
246    path: Union[os.PathLike, str],
247    batch_size: int,
248    patch_shape: Tuple[int, ...],
249    split: Literal["train", "leaderboard", "test"],
250    resize_inputs: bool = False,
251    download: bool = False,
252    **kwargs
253) -> DataLoader:
254    """Get the NCI-ISBI 2013 Prostate dataloader for prostate zone segmentation.
255
256    Args:
257        path: Filepath to a folder where the data is downloaded for further processing.
258        batch_size: The batch size for training.
259        patch_shape: The patch shape to use for training.
260        split: The choice of data split. Either 'train', 'leaderboard' or 'test'.
261        resize_inputs: Whether to resize inputs to the desired patch shape.
262        download: Whether to download the data if it is not present.
263        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
264
265    Returns:
266        The DataLoader.
267    """
268    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
269    dataset = get_nci_isbi_prostate_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
270    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)