torch_em.data.datasets.medical.cc_tumor_heterogeneity

  1"""The CC-Tumor-Heterogeneity (CCTH) dataset contains annotations for cervix and tumor segmentation in T2-weighted
  2MRI of patients with advanced cervical cancer.
  3
  4It consists of 68 T2-weighted MRI volumes (67 sagittal, 1 axial) of 23 patients, acquired at up to three time
  5points during radiochemotherapy. Each volume comes with two contours drawn in MIM: the uterus / cervix (ROI with
  6magenta display color) and the tumor inside of it (ROI with cyan display color), which are distributed as DICOM
  7RTSTRUCT. This module downloads only the RTSTRUCT series and the T2-weighted MRI series they reference (not the
  8DCE / DWI MRI and PET/CT of the collection), rasterizes the contours onto the MRI grid
  9(see `torch_em.data.datasets.util.rasterize_rtstruct`) and stores images and labels in hdf5 files.
 10The semantic label ids are: 1: tumor, 2: cervix / uterus (excluding the tumor).
 11
 12NOTE: This requires the pydicom python package.
 13
 14The dataset is located at https://www.cancerimagingarchive.net/collection/cc-tumor-heterogeneity/.
 15
 16This dataset is from the publication https://doi.org/10.1016/j.ijrobp.2020.02.001.
 17The data was released at https://doi.org/10.7937/TCIA.2022.6X0F-2S2T.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import csv
 23import requests
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from typing import Union, Tuple, List
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36COLLECTION = "CC-Tumor-Heterogeneity"
 37
 38# The DICOM series are downloaded individually from TCIA via the NBIA REST API.
 39URL = util.NBIA_API_URL + "getSeries"
 40CHECKSUM = None
 41
 42LABEL_IDS = {"tumor": 1, "cervix": 2}
 43
 44# The two ROIs of each RTSTRUCT have the same name and are only distinguished by their display color.
 45ROI_COLORS = {(0, 235, 235): LABEL_IDS["tumor"], (255, 0, 255): LABEL_IDS["cervix"]}
 46
 47
 48def _read_rtstruct_info(rtstruct_path):
 49    """Get the referenced series UID and the mapping from ROI numbers to label ids."""
 50    import pydicom
 51
 52    rtstruct = pydicom.dcmread(rtstruct_path, stop_before_pixels=True)
 53    referenced_series = str(
 54        rtstruct.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence[0]
 55        .RTReferencedSeriesSequence[0].SeriesInstanceUID
 56    )
 57    roi_labels = {}
 58    for roi_contour in rtstruct.ROIContourSequence:
 59        color = tuple(int(v) for v in roi_contour.ROIDisplayColor)
 60        assert color in ROI_COLORS, f"Unexpected ROI color {color} in {rtstruct_path}."
 61        roi_labels[int(roi_contour.ReferencedROINumber)] = ROI_COLORS[color]
 62    assert sorted(roi_labels.values()) == sorted(LABEL_IDS.values()), f"Unexpected ROIs in {rtstruct_path}."
 63    roi_names = {int(roi.ROINumber): str(roi.ROIName) for roi in rtstruct.StructureSetROISequence}
 64    return referenced_series, roi_labels, roi_names
 65
 66
 67def _preprocess_cc_tumor_heterogeneity(dicom_dir, csv_path, preprocessed_dir):
 68    import h5py
 69
 70    with open(csv_path, "r") as f:
 71        rtstruct_series = {row["Series UID"]: row["Subject ID"] for row in csv.DictReader(f)}
 72
 73    os.makedirs(preprocessed_dir, exist_ok=True)
 74    for series_uid, subject_id in tqdm(sorted(rtstruct_series.items()), desc="Preprocess CC-Tumor-Heterogeneity"):
 75        rtstruct_path = glob(os.path.join(dicom_dir, series_uid, "*.dcm"))[0]
 76        referenced_series, roi_labels, roi_names = _read_rtstruct_info(rtstruct_path)
 77        # The ROI names encode the imaging plane and time point, e.g. 'Ut-MRT2-Sag-1'.
 78        time_point = roi_names[1].replace("Ut-MRT2-", "")
 79        out_path = os.path.join(preprocessed_dir, f"{subject_id}_{time_point}.h5")
 80        if os.path.exists(out_path):
 81            continue
 82
 83        volume, geometry = util.load_dicom_series(os.path.join(dicom_dir, referenced_series))
 84        labels = util.rasterize_rtstruct(
 85            rtstruct_path, geometry, volume.shape, lambda roi_number, roi_name: roi_labels[roi_number]
 86        )
 87
 88        with h5py.File(out_path, "w") as f:
 89            f.create_dataset("raw", data=volume, compression="gzip")
 90            f.create_dataset("labels", data=labels, compression="gzip")
 91
 92
 93def get_cc_tumor_heterogeneity_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 94    """Download the CC-Tumor-Heterogeneity dataset.
 95
 96    Args:
 97        path: Filepath to a folder where the data is downloaded for further processing.
 98        download: Whether to download the data if it is not present.
 99
100    Returns:
101        Filepath where the preprocessed data is stored.
102    """
103    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
104    preprocessed_dir = os.path.join(path, "preprocessed")
105    dicom_dir = os.path.join(path, "dicom")
106    if not os.path.exists(dicom_dir) and not download:
107        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
108    os.makedirs(path, exist_ok=True)
109
110    # Download the RTSTRUCT series of the collection, and then the T2-weighted MRI series they reference.
111    # The series metadata are written after all series are downloaded, so their presence means it is complete.
112    csv_path = os.path.join(path, "cc_tumor_heterogeneity_rtstruct.csv")
113    image_csv_path = os.path.join(path, "cc_tumor_heterogeneity_images")
114    if not os.path.exists(csv_path) or not os.path.exists(f"{image_csv_path}.csv"):
115        response = requests.get(URL, params={"Collection": COLLECTION, "Modality": "RTSTRUCT"})
116        response.raise_for_status()
117        rtstruct_uids = sorted(series["SeriesInstanceUID"] for series in response.json())
118        csv_path = util.download_tcia_series(rtstruct_uids, dicom_dir, csv_path[:-len(".csv")])
119
120        image_uids = sorted(
121            _read_rtstruct_info(glob(os.path.join(dicom_dir, uid, "*.dcm"))[0])[0] for uid in rtstruct_uids
122        )
123        util.download_tcia_series(image_uids, dicom_dir, image_csv_path)
124
125    _preprocess_cc_tumor_heterogeneity(dicom_dir, csv_path, preprocessed_dir)
126    return preprocessed_dir
127
128
129def get_cc_tumor_heterogeneity_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
130    """Get paths to the CC-Tumor-Heterogeneity data.
131
132    Args:
133        path: Filepath to a folder where the data is downloaded for further processing.
134        download: Whether to download the data if it is not present.
135
136    Returns:
137        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
138    """
139    data_dir = get_cc_tumor_heterogeneity_data(path, download)
140    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
141    return volume_paths
142
143
144def get_cc_tumor_heterogeneity_dataset(
145    path: Union[os.PathLike, str],
146    patch_shape: Tuple[int, ...],
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> Dataset:
151    """Get the CC-Tumor-Heterogeneity dataset for cervix and tumor segmentation in T2-weighted MRI.
152
153    Args:
154        path: Filepath to a folder where the data is downloaded for further processing.
155        patch_shape: The patch shape to use for training.
156        resize_inputs: Whether to resize inputs to the desired patch shape.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
159
160    Returns:
161        The segmentation dataset.
162    """
163    volume_paths = get_cc_tumor_heterogeneity_paths(path, download)
164
165    if resize_inputs:
166        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
167        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
168            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
169        )
170
171    return torch_em.default_segmentation_dataset(
172        raw_paths=volume_paths,
173        raw_key="raw",
174        label_paths=volume_paths,
175        label_key="labels",
176        patch_shape=patch_shape,
177        is_seg_dataset=True,
178        **kwargs
179    )
180
181
182def get_cc_tumor_heterogeneity_loader(
183    path: Union[os.PathLike, str],
184    batch_size: int,
185    patch_shape: Tuple[int, ...],
186    resize_inputs: bool = False,
187    download: bool = False,
188    **kwargs
189) -> DataLoader:
190    """Get the CC-Tumor-Heterogeneity dataloader for cervix and tumor segmentation in T2-weighted MRI.
191
192    Args:
193        path: Filepath to a folder where the data is downloaded for further processing.
194        batch_size: The batch size for training.
195        patch_shape: The patch shape to use for training.
196        resize_inputs: Whether to resize inputs to the desired patch shape.
197        download: Whether to download the data if it is not present.
198        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
199
200    Returns:
201        The DataLoader.
202    """
203    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
204    dataset = get_cc_tumor_heterogeneity_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
205    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)