torch_em.data.datasets.medical.four_d_lung

  1"""The 4D-Lung dataset contains longitudinal respiratory-correlated 4D fan-beam CT (4D-FBCT) and
  24D cone-beam CT (4D-CBCT) acquisitions of 20 locally-advanced non-small cell lung cancer (NSCLC)
  3patients, collected over the course of chemoradiotherapy at the University of Texas MD Anderson
  4Cancer Center. Each patient contributes multiple 4D acquisitions (up to several hundred CT series
  5in total, ~589 studies for the full collection), of which a subset are paired with DICOM RTSTRUCT
  6gross tumor volume (GTV) contours drawn by radiation oncologists.
  7
  8This loader only downloads the CT series that are referenced by an RTSTRUCT (rather than the full,
  9~183 GB collection), and rasterizes the GTV contours onto the CT grid with
 10`torch_em.data.datasets.util.rasterize_rtstruct`, storing image and label pairs as hdf5 files.
 11The labels are binary (foreground: GTV).
 12
 13NOTE: This requires the pydicom python package.
 14
 15The dataset is located at https://doi.org/10.7937/K9/TCIA.2016.ELN8YGLE and is fully public
 16(no data use agreement required), released under a CC-BY-3.0 license.
 17
 18This dataset is from the publication https://doi.org/10.1002/mp.12059.
 19Please cite it if you use this dataset for your research.
 20"""
 21
 22import os
 23import json
 24from glob import glob
 25from natsort import natsorted
 26from typing import Union, Tuple, List
 27
 28import requests
 29from tqdm import tqdm
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38NBIA_API_URL = "https://services.cancerimagingarchive.net/nbia-api/services/v1/"
 39COLLECTION = "4D-Lung"
 40
 41
 42def _get_series(path, modality):
 43    cache_path = os.path.join(path, f"series_{modality.lower()}.json")
 44    if os.path.exists(cache_path):
 45        with open(cache_path, "r") as f:
 46            return json.load(f)
 47
 48    response = requests.get(NBIA_API_URL + "getSeries", params={"Collection": COLLECTION, "Modality": modality})
 49    response.raise_for_status()
 50    series = response.json()
 51
 52    os.makedirs(path, exist_ok=True)
 53    with open(cache_path, "w") as f:
 54        json.dump(series, f)
 55
 56    return series
 57
 58
 59def _referenced_series_uid(rtstruct_path):
 60    import pydicom
 61
 62    rtstruct = pydicom.dcmread(rtstruct_path, stop_before_pixels=True)
 63    referenced_study = rtstruct.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence[0]
 64    return str(referenced_study.RTReferencedSeriesSequence[0].SeriesInstanceUID)
 65
 66
 67def _preprocess_four_d_lung(rtstruct_series, rtstruct_dir, image_dir, preprocessed_dir):
 68    import h5py
 69    import numpy as np
 70
 71    os.makedirs(preprocessed_dir, exist_ok=True)
 72
 73    for series in tqdm(rtstruct_series, desc="Preprocess 4D-Lung"):
 74        uid = series["SeriesInstanceUID"]
 75        patient_id = series["PatientID"]
 76        out_path = os.path.join(preprocessed_dir, f"{patient_id}_{uid[-8:]}.h5")
 77        if os.path.exists(out_path):
 78            continue
 79
 80        rtstruct_path = glob(os.path.join(rtstruct_dir, uid, "*.dcm"))[0]
 81        referenced_uid = _referenced_series_uid(rtstruct_path)
 82        image_series_dir = os.path.join(image_dir, referenced_uid)
 83        if not os.path.exists(image_series_dir):
 84            continue
 85
 86        volume, geometry = util.load_dicom_series(image_series_dir)
 87        volume = np.round(volume).astype("int16")
 88        labels = util.rasterize_rtstruct(rtstruct_path, geometry, volume.shape, lambda roi_number, roi_name: 1)
 89
 90        with h5py.File(out_path, "w") as f:
 91            f.create_dataset("raw", data=volume, compression="gzip")
 92            f.create_dataset("labels", data=labels, compression="gzip")
 93
 94
 95def get_four_d_lung_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 96    """Download the 4D-Lung dataset.
 97
 98    NOTE: This only downloads the CT series that have a matching RTSTRUCT (a small subset of the
 99    full, ~183 GB collection). Ensure enough disk space is available before running this with
100    `download=True`, and expect the initial call to take a while, since every RTSTRUCT of the
101    collection is inspected to find its referenced CT series.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        Filepath where the preprocessed data is stored.
109    """
110    preprocessed_dir = os.path.join(path, "preprocessed")
111    if glob(os.path.join(preprocessed_dir, "*.h5")):
112        return preprocessed_dir
113
114    os.makedirs(path, exist_ok=True)
115    if not download:
116        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
117
118    rtstruct_series = _get_series(path, "RTSTRUCT")
119    rtstruct_dir = os.path.join(path, "rtstructs")
120    util.download_tcia_series(
121        [series["SeriesInstanceUID"] for series in rtstruct_series],
122        dst=rtstruct_dir, csv_filename=os.path.join(path, "rtstructs"),
123    )
124
125    referenced_uids = sorted({
126        _referenced_series_uid(glob(os.path.join(rtstruct_dir, series["SeriesInstanceUID"], "*.dcm"))[0])
127        for series in rtstruct_series
128    })
129    image_dir = os.path.join(path, "images")
130    util.download_tcia_series(referenced_uids, dst=image_dir, csv_filename=os.path.join(path, "images"))
131
132    _preprocess_four_d_lung(rtstruct_series, rtstruct_dir, image_dir, preprocessed_dir)
133    return preprocessed_dir
134
135
136def get_four_d_lung_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
137    """Get paths to the 4D-Lung data.
138
139    Args:
140        path: Filepath to a folder where the data is downloaded for further processing.
141        download: Whether to download the data if it is not present.
142
143    Returns:
144        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
145    """
146    data_dir = get_four_d_lung_data(path, download)
147    return natsorted(glob(os.path.join(data_dir, "*.h5")))
148
149
150def get_four_d_lung_dataset(
151    path: Union[os.PathLike, str],
152    patch_shape: Tuple[int, int, int],
153    resize_inputs: bool = False,
154    download: bool = False,
155    **kwargs
156) -> Dataset:
157    """Get the 4D-Lung dataset for GTV segmentation in 4D lung CT.
158
159    Args:
160        path: Filepath to a folder where the data is downloaded for further processing.
161        patch_shape: The patch shape to use for training.
162        resize_inputs: Whether to resize the inputs to the patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
165
166    Returns:
167        The segmentation dataset.
168    """
169    volume_paths = get_four_d_lung_paths(path, download)
170
171    if resize_inputs:
172        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
173        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
174            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
175        )
176
177    return torch_em.default_segmentation_dataset(
178        raw_paths=volume_paths,
179        raw_key="raw",
180        label_paths=volume_paths,
181        label_key="labels",
182        patch_shape=patch_shape,
183        is_seg_dataset=True,
184        **kwargs
185    )
186
187
188def get_four_d_lung_loader(
189    path: Union[os.PathLike, str],
190    batch_size: int,
191    patch_shape: Tuple[int, int, int],
192    resize_inputs: bool = False,
193    download: bool = False,
194    **kwargs
195) -> DataLoader:
196    """Get the 4D-Lung dataloader for GTV segmentation in 4D lung CT.
197
198    Args:
199        path: Filepath to a folder where the data is downloaded for further processing.
200        batch_size: The batch size for training.
201        patch_shape: The patch shape to use for training.
202        resize_inputs: Whether to resize the inputs to the patch shape.
203        download: Whether to download the data if it is not present.
204        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
205
206    Returns:
207        The DataLoader.
208    """
209    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
210    dataset = get_four_d_lung_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
211    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)