torch_em.data.datasets.medical.nih_pancreas

  1"""The NIH Pancreas-CT dataset contains annotations for pancreas segmentation in contrast-enhanced abdominal CT.
  2
  3It consists of 80 CT volumes (the original release had 82 volumes, cases 25 and 70 were removed by TCIA,
  4as they are duplicates of case 2) with binary pancreas labels. The CT scans are distributed as DICOM series,
  5which are stacked into volumes and stored together with the corresponding labels in hdf5 files by this module.
  6
  7NOTE: This requires the pydicom python package.
  8
  9The dataset is located at https://www.cancerimagingarchive.net/collection/pancreas-ct/.
 10
 11This dataset is from the publication https://doi.org/10.1007/978-3-319-24553-9_68.
 12The data was released at https://doi.org/10.7937/K9/TCIA.2016.tNB1kqBU.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17import csv
 18from glob import glob
 19from tqdm import tqdm
 20from natsort import natsorted
 21from typing import Union, Tuple, List
 22
 23import numpy as np
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URLS = {
 33    "images": "https://www.cancerimagingarchive.net/wp-content/uploads/Pancreas-CT-20200910.tcia",
 34    "labels": "https://www.cancerimagingarchive.net/wp-content/uploads/TCIA_pancreas_labels-02-05-2017-1.zip",
 35}
 36
 37CHECKSUMS = {
 38    "images": None,  # The DICOM series are downloaded individually from TCIA.
 39    "labels": "cf8a553c37c80e3840ce7392987b308f7a98cd5f50019511b2ac0f3a54b0934b",
 40}
 41
 42
 43def _load_dicom_volume(series_dir):
 44    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
 45
 46    Returns the volume in Hounsfield units and the image orientation (DICOM 'ImageOrientationPatient').
 47    """
 48    import pydicom
 49
 50    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 51    slices.sort(key=lambda dcm: float(dcm.ImagePositionPatient[2]))
 52
 53    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 54    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 55    volume = np.round(volume).astype("int16")
 56
 57    orientation = np.round([float(v) for v in slices[0].ImageOrientationPatient]).astype("int").tolist()
 58    return volume, orientation
 59
 60
 61def _preprocess_nih_pancreas(dicom_dir, label_dir, csv_path, preprocessed_dir):
 62    import h5py
 63    import nibabel as nib
 64
 65    with open(csv_path, "r") as f:
 66        subject_ids = {row["Series UID"]: row["Subject ID"] for row in csv.DictReader(f)}
 67
 68    os.makedirs(preprocessed_dir, exist_ok=True)
 69    for series_dir in tqdm(natsorted(glob(os.path.join(dicom_dir, "*"))), desc="Preprocess NIH Pancreas-CT"):
 70        subject_id = subject_ids[os.path.basename(series_dir)]
 71        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
 72        if os.path.exists(out_path):
 73            continue
 74
 75        volume, orientation = _load_dicom_volume(series_dir)
 76        # The volume has axes (z, y, x) with x pointing to the patient's left and y to the anterior (DICOM LPS
 77        # convention with 'ImageOrientationPatient' [1, 0, 0, 0, -1, 0]). The labels are stored with the axis
 78        # orientation (L, A, I), so they only have to be transposed and flipped along z to match the volume.
 79        assert orientation == [1, 0, 0, 0, -1, 0], f"Unexpected image orientation for {subject_id}: {orientation}"
 80
 81        label_path = os.path.join(label_dir, f"label{subject_id.split('_')[-1]}.nii.gz")
 82        label_nifti = nib.load(label_path)
 83        assert nib.aff2axcodes(label_nifti.affine) == ("L", "A", "I"), f"Unexpected label axes for {subject_id}"
 84        labels = np.asarray(label_nifti.dataobj).astype("uint8").transpose(2, 1, 0)[::-1]
 85        assert labels.shape == volume.shape, f"Shape mismatch for {subject_id}: {labels.shape} vs {volume.shape}"
 86
 87        with h5py.File(out_path, "w") as f:
 88            f.create_dataset("raw", data=volume, compression="gzip")
 89            f.create_dataset("labels", data=labels, compression="gzip")
 90
 91
 92def get_nih_pancreas_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 93    """Download the NIH Pancreas-CT dataset.
 94
 95    Args:
 96        path: Filepath to a folder where the data is downloaded for further processing.
 97        download: Whether to download the data if it is not present.
 98
 99    Returns:
100        Filepath where the preprocessed data is stored.
101    """
102    preprocessed_dir = os.path.join(path, "preprocessed")
103    if os.path.exists(preprocessed_dir):
104        return preprocessed_dir
105
106    os.makedirs(path, exist_ok=True)
107
108    # Download the labels.
109    label_dir = os.path.join(path, "TCIA_pancreas_labels-02-05-2017")
110    if not os.path.exists(label_dir):
111        zip_path = os.path.join(path, "TCIA_pancreas_labels-02-05-2017-1.zip")
112        util.download_source(path=zip_path, url=URLS["labels"], download=download, checksum=CHECKSUMS["labels"])
113        util.unzip(zip_path=zip_path, dst=path)
114
115    # Download the DICOM series from the TCIA manifest.
116    dicom_dir = os.path.join(path, "dicom")
117    csv_path = os.path.join(path, "nih_pancreas_series")
118    util.download_source_tcia(
119        path=os.path.join(path, "Pancreas-CT-20200910.tcia"), url=URLS["images"], dst=dicom_dir,
120        csv_filename=csv_path, download=download,
121    )
122
123    _preprocess_nih_pancreas(dicom_dir, label_dir, f"{csv_path}.csv", preprocessed_dir)
124    return preprocessed_dir
125
126
127def get_nih_pancreas_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
128    """Get paths to the NIH Pancreas-CT data.
129
130    Args:
131        path: Filepath to a folder where the data is downloaded for further processing.
132        download: Whether to download the data if it is not present.
133
134    Returns:
135        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
136    """
137    data_dir = get_nih_pancreas_data(path, download)
138    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
139    return volume_paths
140
141
142def get_nih_pancreas_dataset(
143    path: Union[os.PathLike, str],
144    patch_shape: Tuple[int, ...],
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs
148) -> Dataset:
149    """Get the NIH Pancreas-CT dataset for pancreas segmentation.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        patch_shape: The patch shape to use for training.
154        resize_inputs: Whether to resize inputs to the desired patch shape.
155        download: Whether to download the data if it is not present.
156        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
157
158    Returns:
159        The segmentation dataset.
160    """
161    volume_paths = get_nih_pancreas_paths(path, download)
162
163    if resize_inputs:
164        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
165        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
166            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
167        )
168
169    return torch_em.default_segmentation_dataset(
170        raw_paths=volume_paths,
171        raw_key="raw",
172        label_paths=volume_paths,
173        label_key="labels",
174        patch_shape=patch_shape,
175        is_seg_dataset=True,
176        **kwargs
177    )
178
179
180def get_nih_pancreas_loader(
181    path: Union[os.PathLike, str],
182    batch_size: int,
183    patch_shape: Tuple[int, ...],
184    resize_inputs: bool = False,
185    download: bool = False,
186    **kwargs
187) -> DataLoader:
188    """Get the NIH Pancreas-CT dataloader for pancreas segmentation.
189
190    Args:
191        path: Filepath to a folder where the data is downloaded for further processing.
192        batch_size: The batch size for training.
193        patch_shape: The patch shape to use for training.
194        resize_inputs: Whether to resize inputs to the desired patch shape.
195        download: Whether to download the data if it is not present.
196        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
197
198    Returns:
199        The DataLoader.
200    """
201    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
202    dataset = get_nih_pancreas_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
203    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)