torch_em.data.datasets.medical.rider_lung

  1"""The RIDER Lung CT dataset contains annotations for lung tumor segmentation in same-day repeat CT scans
  2of non-small cell lung cancer patients.
  3
  4It consists of 59 CT volumes (test and retest scan of 31 patients, with one scan missing for three patients)
  5with a delineation of the primary gross tumor volume by a radiation oncologist, which is distributed as DICOM
  6RTSTRUCT in the 'RIDER-LungCT-Seg' analysis result of the collection. Each RTSTRUCT contains a manual
  7delineation ('GTVp_<scan>_man') and the result of an in-house autosegmentation method ('GTVp_<scan>_auto').
  8This module rasterizes the contours onto the CT grid (see `torch_em.data.datasets.util.rasterize_rtstruct`)
  9and stores the CT ('raw') and the binary tumor labels ('labels/manual' and 'labels/auto') in hdf5 files.
 10
 11NOTE: This requires the pydicom python package.
 12
 13The dataset is located at https://www.cancerimagingarchive.net/collection/rider-lung-ct/ and the segmentations
 14at https://www.cancerimagingarchive.net/analysis-result/rider-lungct-seg/.
 15
 16This dataset is from the publications https://doi.org/10.1148/radiol.2522081593 (images)
 17and https://doi.org/10.1038/ncomms5006 (segmentations).
 18The data was released at https://doi.org/10.7937/K9/TCIA.2015.U1X8A5NR and https://doi.org/10.7937/tcia.2020.jit9grk8.
 19Please cite them if you use this dataset in your research.
 20"""
 21
 22import os
 23import csv
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from typing import Union, Tuple, List, Literal
 28
 29import numpy as np
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38URLS = {
 39    "images": "https://www.cancerimagingarchive.net/wp-content/uploads/RIDER-Lung-CT-Original-Scans-for-Leonard-Wee-Feb-10-2020-.tcia",  # noqa
 40    "labels": "https://www.cancerimagingarchive.net/wp-content/uploads/RIDER-Lung-CT-RTSTRUCTS-DICOM-SEGS-Leonard-Wee-Feb-10-2020.tcia",  # noqa
 41}
 42
 43# The DICOM series are downloaded individually from TCIA.
 44CHECKSUMS = {"images": None, "labels": None}
 45
 46ANNOTATIONS = ["manual", "auto"]
 47
 48
 49def _get_referenced_series(rtstruct_path):
 50    import pydicom
 51
 52    rtstruct = pydicom.dcmread(rtstruct_path, stop_before_pixels=True)
 53    return str(
 54        rtstruct.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence[0]
 55        .RTReferencedSeriesSequence[0].SeriesInstanceUID
 56    )
 57
 58
 59def _preprocess_rider_lung(dicom_dir, csv_path, preprocessed_dir):
 60    import h5py
 61
 62    with open(csv_path, "r") as f:
 63        rows = list(csv.DictReader(f))
 64    rtstruct_series = {
 65        row["Series UID"]: (row["Subject ID"], row["Series Description"].lower()) for row in rows
 66        if row["Modality"] == "RTSTRUCT"
 67    }
 68
 69    os.makedirs(preprocessed_dir, exist_ok=True)
 70    for series_uid, (subject_id, scan) in tqdm(sorted(rtstruct_series.items()), desc="Preprocess RIDER Lung CT"):
 71        assert scan in ("test", "retest"), f"Unexpected RTSTRUCT series description: {scan}"
 72        out_path = os.path.join(preprocessed_dir, f"{subject_id}_{scan}.h5")
 73        if os.path.exists(out_path):
 74            continue
 75
 76        rtstruct_path = glob(os.path.join(dicom_dir, series_uid, "*.dcm"))[0]
 77        ct_dir = os.path.join(dicom_dir, _get_referenced_series(rtstruct_path))
 78        volume, geometry = util.load_dicom_series(ct_dir)
 79        volume = np.round(volume).astype("int16")
 80
 81        with h5py.File(out_path, "w") as f:
 82            f.create_dataset("raw", data=volume, compression="gzip")
 83            for annotation in ANNOTATIONS:
 84                roi_labels = {f"GTVp_{scan}_{'man' if annotation == 'manual' else 'auto'}": 1}
 85                labels = util.rasterize_rtstruct(rtstruct_path, geometry, volume.shape, roi_labels)
 86                # Not all RTSTRUCTs contain both delineations, e.g. RIDER-2016615262 has no 'auto' retest contour.
 87                if labels.any():
 88                    f.create_dataset(f"labels/{annotation}", data=labels, compression="gzip")
 89
 90
 91def get_rider_lung_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 92    """Download the RIDER Lung CT dataset.
 93
 94    Args:
 95        path: Filepath to a folder where the data is downloaded for further processing.
 96        download: Whether to download the data if it is not present.
 97
 98    Returns:
 99        Filepath where the preprocessed data is stored.
100    """
101    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
102    preprocessed_dir = os.path.join(path, "preprocessed")
103    os.makedirs(path, exist_ok=True)
104
105    # Download the annotated CT series and the RTSTRUCT (and SEG) series from the two TCIA manifests.
106    # The series metadata are written after all series are downloaded, so their presence means it is complete.
107    dicom_dir = os.path.join(path, "dicom")
108    for name, url in URLS.items():
109        csv_filename = os.path.join(path, f"rider_lung_{name}")
110        if not os.path.exists(f"{csv_filename}.csv"):
111            util.download_source_tcia(
112                path=os.path.join(path, os.path.basename(url)), url=url, dst=dicom_dir,
113                csv_filename=csv_filename, download=download,
114            )
115
116    _preprocess_rider_lung(dicom_dir, os.path.join(path, "rider_lung_labels.csv"), preprocessed_dir)
117    return preprocessed_dir
118
119
120def get_rider_lung_paths(
121    path: Union[os.PathLike, str], annotation: Literal["manual", "auto"] = "manual", download: bool = False
122) -> List[str]:
123    """Get paths to the RIDER Lung CT data.
124
125    Args:
126        path: Filepath to a folder where the data is downloaded for further processing.
127        annotation: The tumor delineation to use as labels, either the 'manual' delineation by a radiation
128            oncologist or the result of the 'auto' segmentation method.
129        download: Whether to download the data if it is not present.
130
131    Returns:
132        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data
133        ('labels/manual' and 'labels/auto').
134    """
135    import h5py
136
137    assert annotation in ANNOTATIONS, f"'{annotation}' is not a valid annotation. Choose one of {ANNOTATIONS}."
138    data_dir = get_rider_lung_data(path, download)
139    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
140
141    # A few volumes do not have both delineations, they are filtered out for the respective annotation.
142    def _has_annotation(volume_path):
143        with h5py.File(volume_path, "r") as f:
144            return f"labels/{annotation}" in f
145
146    return [volume_path for volume_path in volume_paths if _has_annotation(volume_path)]
147
148
149def get_rider_lung_dataset(
150    path: Union[os.PathLike, str],
151    patch_shape: Tuple[int, ...],
152    annotation: Literal["manual", "auto"] = "manual",
153    resize_inputs: bool = False,
154    download: bool = False,
155    **kwargs
156) -> Dataset:
157    """Get the RIDER Lung CT dataset for lung tumor segmentation.
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        annotation: The tumor delineation to use as labels, either the 'manual' delineation by a radiation
163            oncologist or the result of the 'auto' segmentation method.
164        resize_inputs: Whether to resize inputs to the desired patch shape.
165        download: Whether to download the data if it is not present.
166        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
167
168    Returns:
169        The segmentation dataset.
170    """
171    volume_paths = get_rider_lung_paths(path, annotation, download)
172
173    if resize_inputs:
174        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
175        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
176            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
177        )
178
179    return torch_em.default_segmentation_dataset(
180        raw_paths=volume_paths,
181        raw_key="raw",
182        label_paths=volume_paths,
183        label_key=f"labels/{annotation}",
184        patch_shape=patch_shape,
185        is_seg_dataset=True,
186        **kwargs
187    )
188
189
190def get_rider_lung_loader(
191    path: Union[os.PathLike, str],
192    batch_size: int,
193    patch_shape: Tuple[int, ...],
194    annotation: Literal["manual", "auto"] = "manual",
195    resize_inputs: bool = False,
196    download: bool = False,
197    **kwargs
198) -> DataLoader:
199    """Get the RIDER Lung CT dataloader for lung tumor segmentation.
200
201    Args:
202        path: Filepath to a folder where the data is downloaded for further processing.
203        batch_size: The batch size for training.
204        patch_shape: The patch shape to use for training.
205        annotation: The tumor delineation to use as labels, either the 'manual' delineation by a radiation
206            oncologist or the result of the 'auto' segmentation method.
207        resize_inputs: Whether to resize inputs to the desired patch shape.
208        download: Whether to download the data if it is not present.
209        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
210
211    Returns:
212        The DataLoader.
213    """
214    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
215    dataset = get_rider_lung_dataset(path, patch_shape, annotation, resize_inputs, download, **ds_kwargs)
216    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)