torch_em.data.datasets.medical.prostate_edge_cases

  1"""The Prostate-Anatomical-Edge-Cases dataset contains annotations for pelvic organ segmentation in
  2radiotherapy planning CT of prostate cancer patients with anatomical edge cases (e.g. hip prostheses,
  3bowel in the pelvis, large or small bladders, or rectal gas).
  4
  5It consists of 131 CT volumes with semantic labels for 1: bladder, 2: prostate, 3: rectum, 4: left femoral head,
  65: right femoral head. The CT scans are distributed as DICOM series and the labels as DICOM RTSTRUCT contours,
  7which are rasterized on the CT voxel grid (with `skimage.draw.polygon`, slice by slice) and stored together with the
  8CT volumes in hdf5 files by this module. The structures are rasterized in the order femoral heads, bladder, rectum,
  9prostate, so the prostate label takes precedence where contours overlap.
 10
 11NOTE: This requires the pydicom python package.
 12
 13The dataset is located at https://www.cancerimagingarchive.net/collection/prostate-anatomical-edge-cases/.
 14
 15This dataset is from the publication https://doi.org/10.1002/mp.16537.
 16The data was released at https://doi.org/10.7937/013R-ZM07.
 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 natsort import natsorted
 25from typing import Union, Tuple, List
 26
 27import numpy as np
 28from skimage.draw import polygon
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URL = "https://www.cancerimagingarchive.net/wp-content/uploads/Prostate-Anatomical-Edge-Cases-May-2023-manifest.tcia"
 38
 39# The DICOM series are downloaded individually from TCIA.
 40CHECKSUM = None
 41
 42LABEL_IDS = {"Bladder": 1, "Prostate": 2, "Rectum": 3, "Femur_Head_L": 4, "Femur_Head_R": 5}
 43
 44# The order in which the structures are rasterized. Later structures overwrite earlier ones where they overlap.
 45RASTERIZATION_ORDER = ["Femur_Head_L", "Femur_Head_R", "Bladder", "Rectum", "Prostate"]
 46
 47
 48def _load_dicom_volume(series_dir):
 49    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted along the slice normal.
 50
 51    Returns the volume in Hounsfield units and the affine matrix that maps voxel indices (z, y, x)
 52    to DICOM patient coordinates.
 53    """
 54    import pydicom
 55
 56    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 57    orientation = np.array([float(v) for v in slices[0].ImageOrientationPatient])
 58    row_dir, col_dir = orientation[:3], orientation[3:]
 59    normal = np.cross(row_dir, col_dir)
 60    slices.sort(key=lambda dcm: np.dot([float(v) for v in dcm.ImagePositionPatient], normal))
 61
 62    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 63    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 64    volume = np.round(volume).astype("int16")
 65
 66    positions = np.array([[float(v) for v in dcm.ImagePositionPatient] for dcm in slices])
 67    spacing = [float(v) for v in slices[0].PixelSpacing]  # The spacing between rows and between columns.
 68    affine = np.eye(4)
 69    affine[:3, 0] = (positions[-1] - positions[0]) / (len(slices) - 1)
 70    affine[:3, 1] = col_dir * spacing[0]
 71    affine[:3, 2] = row_dir * spacing[1]
 72    affine[:3, 3] = positions[0]
 73    return volume, affine
 74
 75
 76def _load_rtstruct(rtstruct_path, shape, inverse_affine):
 77    """Rasterize the closed planar contours of a DICOM RTSTRUCT object on the voxel grid of the reference CT.
 78
 79    The contour points are mapped from patient coordinates to voxel indices via the inverse affine of the CT,
 80    and each contour is filled with `skimage.draw.polygon` in the slice it belongs to.
 81    """
 82    import pydicom
 83
 84    rtstruct = pydicom.dcmread(rtstruct_path)
 85    roi_names = {int(roi.ROINumber): str(roi.ROIName) for roi in rtstruct.StructureSetROISequence}
 86    contour_sequences = {
 87        roi_names[int(roi_contour.ReferencedROINumber)]: roi_contour.ContourSequence
 88        for roi_contour in rtstruct.ROIContourSequence if "ContourSequence" in roi_contour
 89    }
 90
 91    labels = np.zeros(shape, dtype="uint8")
 92    for structure in RASTERIZATION_ORDER:
 93        if structure not in contour_sequences:
 94            continue
 95        for contour in contour_sequences[structure]:
 96            if contour.ContourGeometricType != "CLOSED_PLANAR":
 97                continue
 98            points = np.array([float(v) for v in contour.ContourData]).reshape(-1, 3)
 99            indices = inverse_affine[:3, :3] @ points.T + inverse_affine[:3, 3:]
100            z = int(np.round(indices[0].mean()))
101            assert np.abs(indices[0] - z).max() < 0.1, f"Non-planar contour in {rtstruct_path}."
102            if z < 0 or z >= shape[0]:
103                continue
104            rows, cols = polygon(indices[1], indices[2], shape=shape[1:])
105            labels[z, rows, cols] = LABEL_IDS[structure]
106
107    return labels
108
109
110def _preprocess_prostate_edge_cases(dicom_dir, csv_path, preprocessed_dir):
111    import h5py
112    import pydicom
113
114    with open(csv_path, "r") as f:
115        rtstruct_series = {
116            row["Subject ID"]: row["Series UID"] for row in csv.DictReader(f) if row["Modality"] == "RTSTRUCT"
117        }
118
119    os.makedirs(preprocessed_dir, exist_ok=True)
120    for subject_id, rtstruct_uid in tqdm(
121        sorted(rtstruct_series.items()), desc="Preprocess Prostate-Anatomical-Edge-Cases"
122    ):
123        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
124        if os.path.exists(out_path):
125            continue
126
127        # The structure set references the CT series it was created for.
128        rtstruct_path = glob(os.path.join(dicom_dir, rtstruct_uid, "*.dcm"))[0]
129        rtstruct = pydicom.dcmread(rtstruct_path, stop_before_pixels=True)
130        referenced_study = rtstruct.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence[0]
131        ct_uid = referenced_study.RTReferencedSeriesSequence[0].SeriesInstanceUID
132        volume, affine = _load_dicom_volume(os.path.join(dicom_dir, ct_uid))
133        labels = _load_rtstruct(rtstruct_path, volume.shape, np.linalg.inv(affine))
134
135        with h5py.File(out_path, "w") as f:
136            f.create_dataset("raw", data=volume, compression="gzip")
137            f.create_dataset("labels", data=labels, compression="gzip")
138
139
140def get_prostate_edge_cases_data(path: Union[os.PathLike, str], download: bool = False) -> str:
141    """Download the Prostate-Anatomical-Edge-Cases dataset.
142
143    Args:
144        path: Filepath to a folder where the data is downloaded for further processing.
145        download: Whether to download the data if it is not present.
146
147    Returns:
148        Filepath where the preprocessed data is stored.
149    """
150    preprocessed_dir = os.path.join(path, "preprocessed")
151    if os.path.exists(preprocessed_dir):
152        return preprocessed_dir
153
154    os.makedirs(path, exist_ok=True)
155
156    # Download the DICOM series (CT and RTSTRUCT) from the TCIA manifest.
157    dicom_dir = os.path.join(path, "dicom")
158    csv_path = os.path.join(path, "prostate_edge_cases_series")
159    util.download_source_tcia(
160        path=os.path.join(path, "Prostate-Anatomical-Edge-Cases-May-2023-manifest.tcia"), url=URL, dst=dicom_dir,
161        csv_filename=csv_path, download=download,
162    )
163
164    _preprocess_prostate_edge_cases(dicom_dir, f"{csv_path}.csv", preprocessed_dir)
165    return preprocessed_dir
166
167
168def get_prostate_edge_cases_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
169    """Get paths to the Prostate-Anatomical-Edge-Cases data.
170
171    Args:
172        path: Filepath to a folder where the data is downloaded for further processing.
173        download: Whether to download the data if it is not present.
174
175    Returns:
176        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
177    """
178    data_dir = get_prostate_edge_cases_data(path, download)
179    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
180    return volume_paths
181
182
183def get_prostate_edge_cases_dataset(
184    path: Union[os.PathLike, str],
185    patch_shape: Tuple[int, ...],
186    resize_inputs: bool = False,
187    download: bool = False,
188    **kwargs
189) -> Dataset:
190    """Get the Prostate-Anatomical-Edge-Cases dataset for pelvic organ segmentation.
191
192    Args:
193        path: Filepath to a folder where the data is downloaded for further processing.
194        patch_shape: The patch shape to use for training.
195        resize_inputs: Whether to resize inputs to the desired patch shape.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
198
199    Returns:
200        The segmentation dataset.
201    """
202    volume_paths = get_prostate_edge_cases_paths(path, download)
203
204    if resize_inputs:
205        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
206        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
207            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
208        )
209
210    return torch_em.default_segmentation_dataset(
211        raw_paths=volume_paths,
212        raw_key="raw",
213        label_paths=volume_paths,
214        label_key="labels",
215        patch_shape=patch_shape,
216        is_seg_dataset=True,
217        **kwargs
218    )
219
220
221def get_prostate_edge_cases_loader(
222    path: Union[os.PathLike, str],
223    batch_size: int,
224    patch_shape: Tuple[int, ...],
225    resize_inputs: bool = False,
226    download: bool = False,
227    **kwargs
228) -> DataLoader:
229    """Get the Prostate-Anatomical-Edge-Cases dataloader for pelvic organ segmentation.
230
231    Args:
232        path: Filepath to a folder where the data is downloaded for further processing.
233        batch_size: The batch size for training.
234        patch_shape: The patch shape to use for training.
235        resize_inputs: Whether to resize inputs to the desired patch shape.
236        download: Whether to download the data if it is not present.
237        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
238
239    Returns:
240        The DataLoader.
241    """
242    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
243    dataset = get_prostate_edge_cases_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
244    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)