torch_em.data.datasets.medical.nsclc_radiomics

  1"""The NSCLC-Radiomics dataset (Lung1) contains annotations for the gross tumor volume and thoracic organs
  2in pretreatment CT of non-small cell lung cancer patients.
  3
  4It consists of 422 CT volumes with manual delineations by a radiation oncologist, which are distributed as
  5DICOM RTSTRUCT contours. This module rasterizes the contours onto the CT grid (see
  6`torch_em.data.datasets.util.rasterize_rtstruct`) and stores CT and labels in hdf5 files.
  7The semantic label ids are: 1: tumor, 2: lung ('Lung-Left', 'Lung-Right' or 'Lungs-Total'), 3: heart,
  84: esophagus, 5: spinal cord. The tumor id covers the primary gross tumor volume ('GTV-1', present for all but
  9one patient) and any further 'gtv-*' contours, which are mostly involved lymph nodes and are named
 10inconsistently ('gtv-2', 'gtv_10_r', 'gtv_supraclav', ...). Not all structures are annotated for every patient:
 11of the 422 patients 422 have a tumor, 411 a lung, 355 an esophagus, 411 a spinal cord and 127 a heart contour.
 12The structures to use can be selected via `structures`.
 13
 14NOTE: This requires the pydicom python package.
 15
 16The dataset is located at https://www.cancerimagingarchive.net/collection/nsclc-radiomics/.
 17
 18This dataset is from the publication https://doi.org/10.1038/ncomms5006.
 19The data was released at https://doi.org/10.7937/K9/TCIA.2015.PF0M9REI.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24import csv
 25from glob import glob
 26from tqdm import tqdm
 27from natsort import natsorted
 28from typing import Union, Tuple, List, Optional, Sequence
 29
 30import numpy as np
 31
 32from torch.utils.data import Dataset, DataLoader
 33
 34import torch_em
 35from torch_em.transform.generic import Compose
 36
 37from .. import util
 38
 39
 40URL = "https://www.cancerimagingarchive.net/wp-content/uploads/NSCLC-Radiomics-Version-4-Oct-2020-NBIA-manifest.tcia"
 41
 42# The DICOM series are downloaded individually from TCIA.
 43CHECKSUM = None
 44
 45STRUCTURE_IDS = {"tumor": 1, "lung": 2, "heart": 3, "esophagus": 4, "spinal_cord": 5}
 46
 47
 48class SelectStructures:
 49    """Label transform that keeps only the given label ids and sets all other labels to background.
 50
 51    Args:
 52        label_ids: The label ids to keep.
 53    """
 54    def __init__(self, label_ids: Sequence[int]):
 55        self.label_ids = list(label_ids)
 56
 57    def __call__(self, labels: np.ndarray) -> np.ndarray:
 58        return np.where(np.isin(labels, self.label_ids), labels, 0)
 59
 60
 61def _get_structure_label(roi_number, roi_name):
 62    """Map the (inconsistently named) ROIs of the RTSTRUCT files to the semantic label ids."""
 63    name = roi_name.lower().replace("_", "-").strip()
 64    if name.startswith("gtv"):
 65        return STRUCTURE_IDS["tumor"]
 66    if name.startswith("lung"):
 67        return STRUCTURE_IDS["lung"]
 68    if name.startswith("heart"):
 69        return STRUCTURE_IDS["heart"]
 70    if name.startswith("esophagus"):
 71        return STRUCTURE_IDS["esophagus"]
 72    if name.startswith("spinal-cord") or name.startswith("spinalcord"):
 73        return STRUCTURE_IDS["spinal_cord"]
 74    return None
 75
 76
 77def _get_referenced_series(rtstruct_path):
 78    import pydicom
 79
 80    rtstruct = pydicom.dcmread(rtstruct_path, stop_before_pixels=True)
 81    return str(
 82        rtstruct.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence[0]
 83        .RTReferencedSeriesSequence[0].SeriesInstanceUID
 84    )
 85
 86
 87def _preprocess_nsclc_radiomics(dicom_dir, csv_path, preprocessed_dir):
 88    import h5py
 89
 90    with open(csv_path, "r") as f:
 91        rows = list(csv.DictReader(f))
 92    rtstruct_series = {row["Series UID"]: row["Subject ID"] for row in rows if row["Modality"] == "RTSTRUCT"}
 93
 94    os.makedirs(preprocessed_dir, exist_ok=True)
 95    for series_uid, subject_id in tqdm(sorted(rtstruct_series.items()), desc="Preprocess NSCLC-Radiomics"):
 96        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
 97        if os.path.exists(out_path):
 98            continue
 99
100        rtstruct_path = glob(os.path.join(dicom_dir, series_uid, "*.dcm"))[0]
101        ct_dir = os.path.join(dicom_dir, _get_referenced_series(rtstruct_path))
102        volume, geometry = util.load_dicom_series(ct_dir)
103        volume = np.round(volume).astype("int16")
104        labels = util.rasterize_rtstruct(rtstruct_path, geometry, volume.shape, _get_structure_label)
105
106        with h5py.File(out_path, "w") as f:
107            f.create_dataset("raw", data=volume, compression="gzip")
108            f.create_dataset("labels", data=labels, compression="gzip")
109
110
111def get_nsclc_radiomics_data(path: Union[os.PathLike, str], download: bool = False) -> str:
112    """Download the NSCLC-Radiomics dataset.
113
114    Args:
115        path: Filepath to a folder where the data is downloaded for further processing.
116        download: Whether to download the data if it is not present.
117
118    Returns:
119        Filepath where the preprocessed data is stored.
120    """
121    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
122    preprocessed_dir = os.path.join(path, "preprocessed")
123    os.makedirs(path, exist_ok=True)
124
125    # Download the DICOM series (CT, RTSTRUCT and SEG) from the TCIA manifest. The series metadata are written
126    # after all series are downloaded, so their presence means the download is complete.
127    dicom_dir = os.path.join(path, "dicom")
128    csv_path = os.path.join(path, "nsclc_radiomics_series")
129    if not os.path.exists(f"{csv_path}.csv"):
130        util.download_source_tcia(
131            path=os.path.join(path, os.path.basename(URL)), url=URL, dst=dicom_dir, csv_filename=csv_path,
132            download=download,
133        )
134
135    _preprocess_nsclc_radiomics(dicom_dir, f"{csv_path}.csv", preprocessed_dir)
136    return preprocessed_dir
137
138
139def get_nsclc_radiomics_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
140    """Get paths to the NSCLC-Radiomics data.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        download: Whether to download the data if it is not present.
145
146    Returns:
147        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
148    """
149    data_dir = get_nsclc_radiomics_data(path, download)
150    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
151    return volume_paths
152
153
154def get_nsclc_radiomics_dataset(
155    path: Union[os.PathLike, str],
156    patch_shape: Tuple[int, ...],
157    structures: Optional[Sequence[str]] = None,
158    resize_inputs: bool = False,
159    download: bool = False,
160    **kwargs
161) -> Dataset:
162    """Get the NSCLC-Radiomics dataset for tumor and thoracic organ segmentation in CT.
163
164    Args:
165        path: Filepath to a folder where the data is downloaded for further processing.
166        patch_shape: The patch shape to use for training.
167        structures: The structures to use as labels, a subset of 'tumor', 'lung', 'heart', 'esophagus'
168            and 'spinal_cord'. All other structures are set to background. By default all structures are used.
169        resize_inputs: Whether to resize inputs to the desired patch shape.
170        download: Whether to download the data if it is not present.
171        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
172
173    Returns:
174        The segmentation dataset.
175    """
176    volume_paths = get_nsclc_radiomics_paths(path, download)
177
178    if structures is not None:
179        assert all(structure in STRUCTURE_IDS for structure in structures), f"Invalid structures: {structures}"
180        select_trafo = SelectStructures([STRUCTURE_IDS[structure] for structure in structures])
181        if "label_transform" in kwargs:
182            kwargs["label_transform"] = Compose(select_trafo, kwargs["label_transform"], is_multi_tensor=False)
183        else:
184            kwargs["label_transform"] = select_trafo
185
186    if resize_inputs:
187        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
188        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
189            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
190        )
191
192    return torch_em.default_segmentation_dataset(
193        raw_paths=volume_paths,
194        raw_key="raw",
195        label_paths=volume_paths,
196        label_key="labels",
197        patch_shape=patch_shape,
198        is_seg_dataset=True,
199        **kwargs
200    )
201
202
203def get_nsclc_radiomics_loader(
204    path: Union[os.PathLike, str],
205    batch_size: int,
206    patch_shape: Tuple[int, ...],
207    structures: Optional[Sequence[str]] = None,
208    resize_inputs: bool = False,
209    download: bool = False,
210    **kwargs
211) -> DataLoader:
212    """Get the NSCLC-Radiomics dataloader for tumor and thoracic organ segmentation in CT.
213
214    Args:
215        path: Filepath to a folder where the data is downloaded for further processing.
216        batch_size: The batch size for training.
217        patch_shape: The patch shape to use for training.
218        structures: The structures to use as labels, a subset of 'tumor', 'lung', 'heart', 'esophagus'
219            and 'spinal_cord'. All other structures are set to background. By default all structures are used.
220        resize_inputs: Whether to resize inputs to the desired patch shape.
221        download: Whether to download the data if it is not present.
222        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
223
224    Returns:
225        The DataLoader.
226    """
227    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
228    dataset = get_nsclc_radiomics_dataset(path, patch_shape, structures, resize_inputs, download, **ds_kwargs)
229    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)