torch_em.data.datasets.medical.ct_lymph_nodes

  1"""The CT Lymph Nodes dataset contains annotations for lymph node segmentation in mediastinal and abdominal CT.
  2
  3It consists of 176 CT volumes (90 mediastinal and 86 abdominal scans) with manually traced lymph node
  4segmentations. The labels are instance labels, i.e. each lymph node has its own id. The CT scans are distributed
  5as DICOM series (ca. 58 GB) and the segmentation masks as nifti files, which are stacked, aligned and stored
  6together in hdf5 files by this module.
  7
  8NOTE: This requires the pydicom python package.
  9
 10The dataset is located at https://www.cancerimagingarchive.net/collection/ct-lymph-nodes/.
 11
 12This dataset is from the publications https://doi.org/10.1007/978-3-319-10404-1_65 and
 13https://doi.org/10.1007/978-3-319-24571-3_7 (segmentation masks).
 14The data was released at https://doi.org/10.7937/K9/TCIA.2015.AQIIDCNM.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19import csv
 20from glob import glob
 21from tqdm import tqdm
 22from natsort import natsorted
 23from typing import Union, Tuple, List, Optional, Literal
 24
 25import numpy as np
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URLS = {
 35    "images": "https://www.cancerimagingarchive.net/wp-content/uploads/TCIA_CT_Lymph_Nodes_06-22-2015.tcia",
 36    "labels": "https://www.cancerimagingarchive.net/wp-content/uploads/MED_ABD_LYMPH_MASKS.zip",
 37}
 38
 39CHECKSUMS = {
 40    "images": None,  # The DICOM series are downloaded individually from TCIA.
 41    "labels": "ace3475c21f04c3f3a01e7fa5181fcbcf4a98cc78ea945e058f1b001d25d6745",
 42}
 43
 44REGIONS = {"mediastinal": "MED", "abdominal": "ABD"}
 45
 46
 47def _load_dicom_volume(series_dir):
 48    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
 49
 50    Returns the volume in Hounsfield units and the image orientation (DICOM 'ImageOrientationPatient').
 51    """
 52    import pydicom
 53
 54    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 55    slices.sort(key=lambda dcm: float(dcm.ImagePositionPatient[2]))
 56
 57    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 58    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 59    volume = np.round(volume).astype("int16")
 60
 61    orientation = np.round([float(v) for v in slices[0].ImageOrientationPatient]).astype("int").tolist()
 62    return volume, orientation
 63
 64
 65def _preprocess_ct_lymph_nodes(dicom_dir, label_dir, csv_path, preprocessed_dir):
 66    import h5py
 67    import nibabel as nib
 68
 69    with open(csv_path, "r") as f:
 70        subject_ids = {row["Subject ID"]: row["Series UID"] for row in csv.DictReader(f) if row["Modality"] == "CT"}
 71
 72    os.makedirs(preprocessed_dir, exist_ok=True)
 73    for subject_id, series_uid in tqdm(sorted(subject_ids.items()), desc="Preprocess CT Lymph Nodes"):
 74        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
 75        if os.path.exists(out_path):
 76            continue
 77
 78        volume, orientation = _load_dicom_volume(os.path.join(dicom_dir, series_uid))
 79        # The volume has axes (z, y, x) with x pointing to the patient's left and y to the posterior (DICOM LPS
 80        # convention with 'ImageOrientationPatient' [1, 0, 0, 0, 1, 0]). The labels are stored with the axis
 81        # orientation (L, P, S), so they only have to be transposed to match the volume.
 82        assert orientation == [1, 0, 0, 0, 1, 0], f"Unexpected image orientation for {subject_id}: {orientation}"
 83
 84        label_nifti = nib.load(os.path.join(label_dir, subject_id, f"{subject_id}_mask.nii.gz"))
 85        assert nib.aff2axcodes(label_nifti.affine) == ("L", "P", "S"), f"Unexpected label axes for {subject_id}"
 86        labels = np.asarray(label_nifti.dataobj).astype("uint8").transpose(2, 1, 0)
 87        assert labels.shape == volume.shape, f"Shape mismatch for {subject_id}: {labels.shape} vs {volume.shape}"
 88
 89        with h5py.File(out_path, "w") as f:
 90            f.create_dataset("raw", data=volume, compression="gzip")
 91            f.create_dataset("labels", data=labels, compression="gzip")
 92
 93
 94def get_ct_lymph_nodes_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 95    """Download the CT Lymph Nodes dataset.
 96
 97    Args:
 98        path: Filepath to a folder where the data is downloaded for further processing.
 99        download: Whether to download the data if it is not present.
100
101    Returns:
102        Filepath where the preprocessed data is stored.
103    """
104    preprocessed_dir = os.path.join(path, "preprocessed")
105    if os.path.exists(preprocessed_dir):
106        return preprocessed_dir
107
108    os.makedirs(path, exist_ok=True)
109
110    # Download the labels.
111    label_dir = os.path.join(path, "MED_ABD_LYMPH_MASKS")
112    if not os.path.exists(label_dir):
113        zip_path = os.path.join(path, "MED_ABD_LYMPH_MASKS.zip")
114        util.download_source(path=zip_path, url=URLS["labels"], download=download, checksum=CHECKSUMS["labels"])
115        util.unzip(zip_path=zip_path, dst=path)
116
117    # Download the DICOM series from the TCIA manifest.
118    dicom_dir = os.path.join(path, "dicom")
119    csv_path = os.path.join(path, "ct_lymph_nodes_series")
120    util.download_source_tcia(
121        path=os.path.join(path, "TCIA_CT_Lymph_Nodes_06-22-2015.tcia"), url=URLS["images"], dst=dicom_dir,
122        csv_filename=csv_path, download=download,
123    )
124
125    _preprocess_ct_lymph_nodes(dicom_dir, label_dir, f"{csv_path}.csv", preprocessed_dir)
126    return preprocessed_dir
127
128
129def get_ct_lymph_nodes_paths(
130    path: Union[os.PathLike, str],
131    region: Optional[Literal["mediastinal", "abdominal"]] = None,
132    download: bool = False,
133) -> List[str]:
134    """Get paths to the CT Lymph Nodes data.
135
136    Args:
137        path: Filepath to a folder where the data is downloaded for further processing.
138        region: The choice of body region. Either 'mediastinal' or 'abdominal'. If None, all volumes are returned.
139        download: Whether to download the data if it is not present.
140
141    Returns:
142        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
143    """
144    data_dir = get_ct_lymph_nodes_data(path, download)
145
146    if region is None:
147        prefix = "*"
148    elif region in REGIONS:
149        prefix = REGIONS[region]
150    else:
151        raise ValueError(f"'{region}' is not a valid region. Please choose one of {list(REGIONS.keys())}.")
152
153    volume_paths = natsorted(glob(os.path.join(data_dir, f"{prefix}_LYMPH_*.h5")))
154    return volume_paths
155
156
157def get_ct_lymph_nodes_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, ...],
160    region: Optional[Literal["mediastinal", "abdominal"]] = None,
161    resize_inputs: bool = False,
162    download: bool = False,
163    **kwargs
164) -> Dataset:
165    """Get the CT Lymph Nodes dataset for lymph node segmentation.
166
167    Args:
168        path: Filepath to a folder where the data is downloaded for further processing.
169        patch_shape: The patch shape to use for training.
170        region: The choice of body region. Either 'mediastinal' or 'abdominal'. If None, all volumes are returned.
171        resize_inputs: Whether to resize inputs to the desired patch shape.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
174
175    Returns:
176        The segmentation dataset.
177    """
178    volume_paths = get_ct_lymph_nodes_paths(path, region, download)
179
180    if resize_inputs:
181        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
182        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
183            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
184        )
185
186    return torch_em.default_segmentation_dataset(
187        raw_paths=volume_paths,
188        raw_key="raw",
189        label_paths=volume_paths,
190        label_key="labels",
191        patch_shape=patch_shape,
192        is_seg_dataset=True,
193        **kwargs
194    )
195
196
197def get_ct_lymph_nodes_loader(
198    path: Union[os.PathLike, str],
199    batch_size: int,
200    patch_shape: Tuple[int, ...],
201    region: Optional[Literal["mediastinal", "abdominal"]] = None,
202    resize_inputs: bool = False,
203    download: bool = False,
204    **kwargs
205) -> DataLoader:
206    """Get the CT Lymph Nodes dataloader for lymph node segmentation.
207
208    Args:
209        path: Filepath to a folder where the data is downloaded for further processing.
210        batch_size: The batch size for training.
211        patch_shape: The patch shape to use for training.
212        region: The choice of body region. Either 'mediastinal' or 'abdominal'. If None, all volumes are returned.
213        resize_inputs: Whether to resize inputs to the desired patch shape.
214        download: Whether to download the data if it is not present.
215        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
216
217    Returns:
218        The DataLoader.
219    """
220    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
221    dataset = get_ct_lymph_nodes_dataset(path, patch_shape, region, resize_inputs, download, **ds_kwargs)
222    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)