torch_em.data.datasets.medical.nsclc_radiogenomics

  1"""The NSCLC-Radiogenomics dataset contains annotations for primary tumor segmentation in preoperative
  2CT of non-small cell lung cancer patients.
  3
  4It consists of CT volumes with a manual delineation of the primary gross tumor volume by a radiation
  5oncologist, distributed as DICOM-SEG objects, which are converted and stored in hdf5 files by this
  6module. The semantic label id is: 1: tumor. Only 144 of the 211 patients in the collection have a
  7segmentation; this module queries the TCIA REST API for the segmented series and downloads only those
  8and their referenced CT series, rather than the full ~98GB collection (which also has PET series and
  9CT series for the unsegmented patients).
 10
 11NOTE: This requires the pydicom python package.
 12
 13The dataset is located at https://www.cancerimagingarchive.net/collection/nsclc-radiogenomics/.
 14
 15This dataset is from the publication https://doi.org/10.1038/sdata.2018.202.
 16The data was released at https://doi.org/10.7937/K9/TCIA.2017.7hs46erv.
 17Please cite it if you use this dataset in your research.
 18"""
 19
 20import os
 21import csv
 22from glob import glob
 23from tqdm import tqdm
 24from warnings import warn
 25from natsort import natsorted
 26from collections import defaultdict, Counter
 27from typing import Union, Tuple, List
 28
 29import requests
 30import numpy as np
 31
 32from torch.utils.data import Dataset, DataLoader
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39COLLECTION = "NSCLC Radiogenomics"
 40
 41LABEL_IDS = {"tumor": 1}
 42
 43
 44def _get_segmented_series_uids():
 45    """Query the TCIA REST API for the SEG series of the collection, without downloading any image data."""
 46    response = requests.get(util.NBIA_API_URL + "getSeries", params={"Collection": COLLECTION, "Modality": "SEG"})
 47    response.raise_for_status()
 48    return [row["SeriesInstanceUID"] for row in response.json()]
 49
 50
 51def _get_referenced_sop_uids(seg_path):
 52    """Get the SOP instance UIDs of the CT slices referenced by a DICOM-SEG object."""
 53    import pydicom
 54
 55    seg = pydicom.dcmread(seg_path, stop_before_pixels=True)
 56    assert len(seg.ReferencedSeriesSequence) == 1, f"Expected a single referenced CT series in {seg_path}."
 57    referenced_instances = seg.ReferencedSeriesSequence[0].ReferencedInstanceSequence
 58    return {str(instance.ReferencedSOPInstanceUID) for instance in referenced_instances}
 59
 60
 61def _get_referenced_ct_series_uid(seg_path):
 62    """Get the series instance UID of the CT series referenced by a DICOM-SEG object."""
 63    import pydicom
 64
 65    seg = pydicom.dcmread(seg_path, stop_before_pixels=True)
 66    assert len(seg.ReferencedSeriesSequence) == 1, f"Expected a single referenced CT series in {seg_path}."
 67    return str(seg.ReferencedSeriesSequence[0].SeriesInstanceUID)
 68
 69
 70def _load_dicom_volume(series_dir, referenced_sop_uids):
 71    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
 72
 73    Some series contain several acquisitions with overlapping slice positions. In this case, only the
 74    acquisition referenced by the DICOM-SEG object is kept, so that each slice position occurs once.
 75    Returns the volume in Hounsfield units and the geometry needed to align the DICOM-SEG frames with the volume.
 76    """
 77    import pydicom
 78
 79    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 80
 81    acquisitions = Counter(dcm.get("AcquisitionNumber") for dcm in slices if dcm.SOPInstanceUID in referenced_sop_uids)
 82    acquisition = acquisitions.most_common(1)[0][0] if acquisitions else None
 83    slices_per_position = {}
 84    for dcm in slices:
 85        z = round(float(dcm.ImagePositionPatient[2]), 2)
 86        priority = (dcm.SOPInstanceUID in referenced_sop_uids, dcm.get("AcquisitionNumber") == acquisition)
 87        if z not in slices_per_position or priority > slices_per_position[z][0]:
 88            slices_per_position[z] = (priority, dcm)
 89    slices = [dcm for _, dcm in sorted(slices_per_position.values(), key=lambda item: item[1].ImagePositionPatient[2])]
 90
 91    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 92    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 93    volume = np.round(volume).astype("int16")
 94
 95    geometry = {
 96        "sop_uids": {str(dcm.SOPInstanceUID): i for i, dcm in enumerate(slices)},
 97        "z_positions": np.array([float(dcm.ImagePositionPatient[2]) for dcm in slices]),
 98        "orientation": np.round([float(v) for v in slices[0].ImageOrientationPatient]).astype("int"),
 99    }
100    return volume, geometry
101
102
103def _load_dicom_seg(seg_path, shape, geometry):
104    """Convert a DICOM-SEG object into a binary tumor mask aligned with the reference CT volume.
105
106    Each frame is mapped to its CT slice via the source image it was derived from (or its z position).
107    All segments in the object are combined into a single tumor mask, since the collection ships one
108    primary gross tumor volume per patient.
109    """
110    import pydicom
111
112    seg = pydicom.dcmread(seg_path)
113    frames = seg.pixel_array
114    if frames.ndim == 2:  # A segmentation with a single frame.
115        frames = frames[None]
116
117    # The segmentation frames may use a different in-plane orientation than the CT slices,
118    # in which case they have to be flipped to align them.
119    seg_orientation = seg.SharedFunctionalGroupsSequence[0].PlaneOrientationSequence[0].ImageOrientationPatient
120    seg_orientation = np.round([float(v) for v in seg_orientation]).astype("int")
121    orientation = geometry["orientation"]
122    assert np.all(np.abs(seg_orientation) == np.abs(orientation)), f"Unexpected orientation in {seg_path}."
123    flip_axes = []
124    if np.any(seg_orientation[3:] != orientation[3:]):  # The direction of the rows differs.
125        flip_axes.append(0)
126    if np.any(seg_orientation[:3] != orientation[:3]):  # The direction of the columns differs.
127        flip_axes.append(1)
128
129    # Frames are matched to CT slices via their z position, if they cannot be matched via the source image.
130    z_positions = geometry["z_positions"]
131    tolerance = np.diff(z_positions).min() / 2 if len(z_positions) > 1 else 1.0
132
133    mask = np.zeros(shape, dtype="bool")
134    for frame, frame_group in zip(frames, seg.PerFrameFunctionalGroupsSequence):
135        frame_mask = frame.astype("bool")
136
137        z = None
138        derivation = frame_group.get("DerivationImageSequence", [])
139        if derivation and derivation[0].get("SourceImageSequence"):
140            z = geometry["sop_uids"].get(str(derivation[0].SourceImageSequence[0].ReferencedSOPInstanceUID))
141        if z is None:
142            frame_z = float(frame_group.PlanePositionSequence[0].ImagePositionPatient[2])
143            z = int(np.argmin(np.abs(z_positions - frame_z)))
144            if abs(z_positions[z] - frame_z) > tolerance:
145                if frame_mask.any():
146                    warn(f"Skipping a frame at z={frame_z} in {seg_path}, which does not match a CT slice.")
147                continue
148
149        if flip_axes:
150            frame_mask = np.flip(frame_mask, axis=flip_axes)
151        mask[z] |= frame_mask
152
153    return mask
154
155
156def _preprocess_nsclc_radiogenomics(dicom_dir, csv_paths, preprocessed_dir):
157    import h5py
158
159    series_per_subject = defaultdict(dict)
160    for csv_path in csv_paths:
161        with open(csv_path, "r") as f:
162            for row in csv.DictReader(f):
163                series_per_subject[row["Subject ID"]][row["Modality"]] = os.path.join(dicom_dir, row["Series UID"])
164
165    os.makedirs(preprocessed_dir, exist_ok=True)
166    subjects_with_seg = {sid: series for sid, series in series_per_subject.items() if "SEG" in series}
167    for subject_id, series_dirs in tqdm(sorted(subjects_with_seg.items()), desc="Preprocess NSCLC-Radiogenomics"):
168        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
169        if os.path.exists(out_path):
170            continue
171        if "CT" not in series_dirs:
172            warn(f"Skipping {subject_id}, which has a SEG object but no CT series.")
173            continue
174
175        seg_path = glob(os.path.join(series_dirs["SEG"], "*.dcm"))[0]
176        volume, geometry = _load_dicom_volume(series_dirs["CT"], _get_referenced_sop_uids(seg_path))
177        mask = _load_dicom_seg(seg_path, volume.shape, geometry)
178        labels = (mask * LABEL_IDS["tumor"]).astype("uint8")
179
180        with h5py.File(out_path, "w") as f:
181            f.create_dataset("raw", data=volume, compression="gzip")
182            f.create_dataset("labels", data=labels, compression="gzip")
183
184
185def get_nsclc_radiogenomics_data(path: Union[os.PathLike, str], download: bool = False) -> str:
186    """Download the NSCLC-Radiogenomics dataset.
187
188    Args:
189        path: Filepath to a folder where the data is downloaded for further processing.
190        download: Whether to download the data if it is not present.
191
192    Returns:
193        Filepath where the preprocessed data is stored.
194    """
195    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
196    preprocessed_dir = os.path.join(path, "preprocessed")
197    os.makedirs(path, exist_ok=True)
198
199    # Only the segmented series and their referenced CT series are downloaded (see the module docstring),
200    # not the full collection manifest. Each download step is skipped once its metadata csv is written.
201    dicom_dir = os.path.join(path, "dicom")
202    seg_csv_path = os.path.join(path, "nsclc_radiogenomics_seg_series")
203    if not os.path.exists(f"{seg_csv_path}.csv"):
204        if not download:
205            raise RuntimeError(f"Cannot find the data at {dicom_dir}, but download was set to False.")
206        seg_uids = _get_segmented_series_uids()
207        util.download_tcia_series(seg_uids, dicom_dir, seg_csv_path)
208
209    ct_csv_path = os.path.join(path, "nsclc_radiogenomics_ct_series")
210    if not os.path.exists(f"{ct_csv_path}.csv"):
211        with open(f"{seg_csv_path}.csv", "r") as f:
212            seg_uids = [row["Series UID"] for row in csv.DictReader(f)]
213        ct_uids = sorted({
214            _get_referenced_ct_series_uid(glob(os.path.join(dicom_dir, uid, "*.dcm"))[0]) for uid in seg_uids
215        })
216        util.download_tcia_series(ct_uids, dicom_dir, ct_csv_path)
217
218    _preprocess_nsclc_radiogenomics(dicom_dir, [f"{seg_csv_path}.csv", f"{ct_csv_path}.csv"], preprocessed_dir)
219    return preprocessed_dir
220
221
222def get_nsclc_radiogenomics_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
223    """Get paths to the NSCLC-Radiogenomics data.
224
225    Args:
226        path: Filepath to a folder where the data is downloaded for further processing.
227        download: Whether to download the data if it is not present.
228
229    Returns:
230        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
231    """
232    data_dir = get_nsclc_radiogenomics_data(path, download)
233    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
234    return volume_paths
235
236
237def get_nsclc_radiogenomics_dataset(
238    path: Union[os.PathLike, str],
239    patch_shape: Tuple[int, ...],
240    resize_inputs: bool = False,
241    download: bool = False,
242    **kwargs
243) -> Dataset:
244    """Get the NSCLC-Radiogenomics dataset for lung tumor segmentation.
245
246    Args:
247        path: Filepath to a folder where the data is downloaded for further processing.
248        patch_shape: The patch shape to use for training.
249        resize_inputs: Whether to resize inputs to the desired patch shape.
250        download: Whether to download the data if it is not present.
251        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
252
253    Returns:
254        The segmentation dataset.
255    """
256    volume_paths = get_nsclc_radiogenomics_paths(path, download)
257
258    if resize_inputs:
259        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
260        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
261            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
262        )
263
264    return torch_em.default_segmentation_dataset(
265        raw_paths=volume_paths,
266        raw_key="raw",
267        label_paths=volume_paths,
268        label_key="labels",
269        patch_shape=patch_shape,
270        is_seg_dataset=True,
271        **kwargs
272    )
273
274
275def get_nsclc_radiogenomics_loader(
276    path: Union[os.PathLike, str],
277    batch_size: int,
278    patch_shape: Tuple[int, ...],
279    resize_inputs: bool = False,
280    download: bool = False,
281    **kwargs
282) -> DataLoader:
283    """Get the NSCLC-Radiogenomics dataloader for lung tumor segmentation.
284
285    Args:
286        path: Filepath to a folder where the data is downloaded for further processing.
287        batch_size: The batch size for training.
288        patch_shape: The patch shape to use for training.
289        resize_inputs: Whether to resize inputs to the desired patch shape.
290        download: Whether to download the data if it is not present.
291        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
292
293    Returns:
294        The DataLoader.
295    """
296    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
297    dataset = get_nsclc_radiogenomics_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
298    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)