torch_em.data.datasets.medical.picai

  1"""The PI-CAI dataset contains annotations for clinically significant prostate cancer and the prostate
  2gland in biparametric MRI.
  3
  4The dataset consists of 1500 biparametric MRI scans of 1476 patients from three centers. Three sets of
  5annotations are provided and selected with the 'annotation' argument: 'lesion_expert' are the csPCa
  6lesions delineated by human experts for 1295 scans, with the ISUP grade as the label id, 'lesion_ai' are
  7the AI derived csPCa lesions for all 1500 scans, with binary labels, and 'whole_gland' is the AI derived
  8prostate mask for all 1500 scans. See also `ANNOTATIONS`.
  9
 10NOTE: Only the axial T2 weighted scan is used. The annotations are resampled to its grid, while the
 11diffusion weighted scans of the same study are acquired on a much coarser grid and do not align with
 12them.
 13
 14NOTE: This requires the SimpleITK python package to read the MetaImage (.mha) scans.
 15
 16The images are located at https://doi.org/10.5281/zenodo.6624726 and the annotations at
 17https://github.com/DIAGNijmegen/picai_labels. Both are distributed under the CC BY-NC 4.0 license.
 18This dataset is from the publication https://doi.org/10.1016/S1470-2045(24)00220-1.
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23from glob import glob
 24from tqdm import tqdm
 25from natsort import natsorted
 26from typing import Union, Tuple, Literal, List
 27
 28import numpy as np
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URL = "https://zenodo.org/records/6624726/files/picai_public_images_fold{fold}.zip?download=1"
 38
 39CHECKSUMS = {
 40    0: "1c9683b436bedbe4384bb4a8133ebb4213349e5af2296ba62cd26fafa36366e0",
 41    1: "eb31d43a949e9245385cea5e66f4584197bbdb7df4ee99af6e47da980ddbb565",
 42    2: "5116355bc0cbf152122467405a0e9636c268e46a144f0f5de51abb48753e9d29",
 43    3: "070f172b5668e9a07900d9f07a48e0a313067e63280140ca0c43a6070487e765",
 44    4: "feabdbf8bb2d28c091dc41fcd66cbd4c2a898aaee05ca27401ccfc343ccec016",
 45}
 46
 47LABELS_URL = "https://github.com/DIAGNijmegen/picai_labels/archive/refs/heads/main.zip"
 48
 49ANNOTATIONS = {
 50    "lesion_expert": "csPCa_lesion_delineations/human_expert/resampled",
 51    "lesion_ai": "csPCa_lesion_delineations/AI/Bosma22a",
 52    "whole_gland": "anatomical_delineations/whole_gland/AI/Bosma22b",
 53}
 54"""Mapping from the annotation choice to its folder in the annotation release."""
 55
 56FOLDS = (0, 1, 2, 3, 4)
 57
 58
 59def _preprocess_picai(image_dir, label_root, preprocessed_dir):
 60    import h5py
 61    import nibabel as nib
 62    import SimpleITK as sitk
 63
 64    os.makedirs(preprocessed_dir, exist_ok=True)
 65    image_paths = natsorted(glob(os.path.join(image_dir, "*", "*_t2w.mha")))
 66    for image_path in tqdm(image_paths, desc="Preprocess PI-CAI"):
 67        case_id = os.path.basename(image_path)[:-len("_t2w.mha")]
 68        out_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 69        if os.path.exists(out_path):
 70            continue
 71
 72        # SimpleITK returns the volume with axis order (z, y, x).
 73        volume = sitk.GetArrayFromImage(sitk.ReadImage(image_path))
 74        with h5py.File(out_path, "w") as f:
 75            f.create_dataset("raw", data=volume, compression="gzip")
 76            for name, rel_path in ANNOTATIONS.items():
 77                label_path = os.path.join(label_root, rel_path, f"{case_id}.nii.gz")
 78                if not os.path.exists(label_path):
 79                    continue
 80                # The annotations are stored as nifti with axis order (x, y, z).
 81                labels = np.asarray(nib.load(label_path).dataobj).T
 82                assert labels.shape == volume.shape, f"The '{name}' mask of '{case_id}' does not match its scan."
 83                f.create_dataset(f"labels/{name}", data=labels.astype("uint8"), compression="gzip")
 84
 85
 86def get_picai_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 87    """Download the PI-CAI dataset.
 88
 89    Args:
 90        path: Filepath to a folder where the data is downloaded for further processing.
 91        download: Whether to download the data if it is not present.
 92
 93    Returns:
 94        Filepath where the preprocessed data is stored.
 95    """
 96    preprocessed_dir = os.path.join(path, "preprocessed")
 97    if os.path.exists(preprocessed_dir) and glob(os.path.join(preprocessed_dir, "*.h5")):
 98        return preprocessed_dir
 99
100    os.makedirs(path, exist_ok=True)
101
102    image_dir = os.path.join(path, "images")
103    for fold in FOLDS:
104        zip_path = os.path.join(path, f"picai_public_images_fold{fold}.zip")
105        util.download_source(
106            path=zip_path, url=URL.format(fold=fold), download=download, checksum=CHECKSUMS[fold]
107        )
108        util.unzip(zip_path=zip_path, dst=image_dir, remove=False)
109
110    label_root = os.path.join(path, "picai_labels-main")
111    if not os.path.exists(label_root):
112        zip_path = os.path.join(path, "picai_labels.zip")
113        util.download_source(path=zip_path, url=LABELS_URL, download=download, checksum=None)
114        util.unzip(zip_path=zip_path, dst=path, remove=False)
115
116    _preprocess_picai(image_dir, label_root, preprocessed_dir)
117    return preprocessed_dir
118
119
120def get_picai_paths(
121    path: Union[os.PathLike, str],
122    annotation: Literal["lesion_expert", "lesion_ai", "whole_gland"] = "lesion_ai",
123    download: bool = False,
124) -> List[str]:
125    """Get paths to the PI-CAI data.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        annotation: The choice of annotations. Either 'lesion_expert', 'lesion_ai' or 'whole_gland'.
130        download: Whether to download the data if it is not present.
131
132    Returns:
133        List of filepaths for the stored data.
134    """
135    import h5py
136
137    if annotation not in ANNOTATIONS:
138        raise ValueError(f"'{annotation}' is not a valid annotation. Choose from {list(ANNOTATIONS.keys())}.")
139
140    preprocessed_dir = get_picai_data(path, download)
141
142    # The human expert lesions are only delineated for a subset of the scans.
143    volume_paths = []
144    for volume_path in natsorted(glob(os.path.join(preprocessed_dir, "*.h5"))):
145        with h5py.File(volume_path, "r") as f:
146            if f"labels/{annotation}" in f:
147                volume_paths.append(volume_path)
148
149    assert len(volume_paths) > 0, f"Could not find any volume with '{annotation}' labels in '{preprocessed_dir}'."
150    return volume_paths
151
152
153def get_picai_dataset(
154    path: Union[os.PathLike, str],
155    patch_shape: Tuple[int, ...],
156    annotation: Literal["lesion_expert", "lesion_ai", "whole_gland"] = "lesion_ai",
157    resize_inputs: bool = False,
158    download: bool = False,
159    **kwargs
160) -> Dataset:
161    """Get the PI-CAI dataset for prostate cancer and prostate gland segmentation.
162
163    Args:
164        path: Filepath to a folder where the data is downloaded for further processing.
165        patch_shape: The patch shape to use for training.
166        annotation: The choice of annotations. Either 'lesion_expert', 'lesion_ai' or 'whole_gland'.
167        resize_inputs: Whether to resize inputs to the desired patch shape.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
170
171    Returns:
172        The segmentation dataset.
173    """
174    volume_paths = get_picai_paths(path, annotation, download)
175
176    if resize_inputs:
177        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
178        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
179            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
180        )
181
182    return torch_em.default_segmentation_dataset(
183        raw_paths=volume_paths,
184        raw_key="raw",
185        label_paths=volume_paths,
186        label_key=f"labels/{annotation}",
187        patch_shape=patch_shape,
188        is_seg_dataset=True,
189        **kwargs
190    )
191
192
193def get_picai_loader(
194    path: Union[os.PathLike, str],
195    batch_size: int,
196    patch_shape: Tuple[int, ...],
197    annotation: Literal["lesion_expert", "lesion_ai", "whole_gland"] = "lesion_ai",
198    resize_inputs: bool = False,
199    download: bool = False,
200    **kwargs
201) -> DataLoader:
202    """Get the PI-CAI dataloader for prostate cancer and prostate gland segmentation.
203
204    Args:
205        path: Filepath to a folder where the data is downloaded for further processing.
206        batch_size: The batch size for training.
207        patch_shape: The patch shape to use for training.
208        annotation: The choice of annotations. Either 'lesion_expert', 'lesion_ai' or 'whole_gland'.
209        resize_inputs: Whether to resize inputs to the desired patch shape.
210        download: Whether to download the data if it is not present.
211        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
212
213    Returns:
214        The DataLoader.
215    """
216    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
217    dataset = get_picai_dataset(path, patch_shape, annotation, resize_inputs, download, **ds_kwargs)
218    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)