torch_em.data.datasets.medical.pediatric_ct_seg

  1"""The Pediatric-CT-SEG dataset contains annotations for organ segmentation in pediatric CT.
  2
  3It consists of 359 CT volumes (chest-abdomen-pelvis, patients aged 5 days to 16 years) with contours of up to
  429 organs, which are distributed as DICOM RTSTRUCT objects. The CT scans are stacked into volumes, the contours are
  5rasterized on the CT grid and both are stored together in hdf5 files by this module.
  6The organs are mapped to the following semantic ids:
  71: adrenal left, 2: adrenal right, 3: bladder, 4: breast left, 5: breast right, 6: duodenum, 7: esophagus,
  88: femoral head left, 9: femoral head right, 10: gall bladder, 11: gonads, 12: heart, 13: kidney left,
  914: kidney right, 15: large intestine, 16: liver, 17: pancreas, 18: prostate, 19: rectum, 20: small intestine,
 1021: spinal canal, 22: spleen, 23: stomach, 24: thymus, 25: uterocervix, 26: lung left, 27: lung right, 28: skin,
 1129: bones.
 12Some volumes are missing structures that are outside of the scan range or could not be identified reliably.
 13The contours of overlapping structures are rasterized in a fixed order, so that the skin and bone labels are
 14overwritten by the organs. Additional contours that do not belong to the 29 organs (e.g. 'BODY' or a
 15'Horseshoe Kidney' in one patient) are not included in the labels.
 16
 17NOTE: This requires the pydicom and opencv python packages.
 18
 19The dataset is located at https://www.cancerimagingarchive.net/collection/pediatric-ct-seg/.
 20
 21This dataset is from the publication https://doi.org/10.1002/mp.15485.
 22The data was released at https://doi.org/10.7937/TCIA.X0H0-1706.
 23Please cite it if you use this dataset in your research.
 24"""
 25
 26import os
 27from glob import glob
 28from tqdm import tqdm
 29from warnings import warn
 30from natsort import natsorted
 31from collections import defaultdict
 32from typing import Union, Tuple, List
 33
 34import numpy as np
 35
 36from torch.utils.data import Dataset, DataLoader
 37
 38import torch_em
 39
 40from .. import util
 41
 42
 43URL = "https://www.cancerimagingarchive.net/wp-content/uploads/Pediatric-CT-SEG-Mar-22-2022-manifest.tcia"
 44
 45# The DICOM series are downloaded individually from TCIA.
 46CHECKSUM = None
 47
 48NUM_VOLUMES = 359
 49
 50# The organ ids follow the ROI numbering of the RTSTRUCT files.
 51ORGAN_IDS = {
 52    "adrenal_left": 1, "adrenal_right": 2, "bladder": 3, "breast_left": 4, "breast_right": 5, "duodenum": 6,
 53    "esophagus": 7, "femoral_head_left": 8, "femoral_head_right": 9, "gall_bladder": 10, "gonads": 11, "heart": 12,
 54    "kidney_left": 13, "kidney_right": 14, "large_intestine": 15, "liver": 16, "pancreas": 17, "prostate": 18,
 55    "rectum": 19, "small_intestine": 20, "spinal_canal": 21, "spleen": 22, "stomach": 23, "thymus": 24,
 56    "uterocervix": 25, "lung_left": 26, "lung_right": 27, "skin": 28, "bones": 29,
 57}
 58
 59# The ROI names used in the RTSTRUCT files (including the alternative spellings in a few patients).
 60ROI_NAMES = {
 61    "Adrenal Left": "adrenal_left", "Lt Adrenal": "adrenal_left", "Adrenal Right": "adrenal_right",
 62    "Rt Adrenal": "adrenal_right", "Bladder": "bladder", "Breast Left": "breast_left", "Breast Right": "breast_right",
 63    "Duodenum": "duodenum", "Esophagus": "esophagus", "Femoral Head Lef": "femoral_head_left",
 64    "Femoral Head Rig": "femoral_head_right", "Gall Bladder": "gall_bladder", "Gonads": "gonads", "Heart": "heart",
 65    "Kidney Left": "kidney_left", "Kidney Right": "kidney_right", "Large Intestine": "large_intestine",
 66    "Liver": "liver", "Pancreas": "pancreas", "Prostate": "prostate", "Rectum": "rectum",
 67    "Small Intestine": "small_intestine", "Spinal Canal": "spinal_canal", "Spleen": "spleen", "Stomach": "stomach",
 68    "Thymus": "thymus", "UteroCervix": "uterocervix", "Lung_L": "lung_left", "Lung_R": "lung_right", "Skin": "skin",
 69    "Bones": "bones",
 70}
 71
 72# The order in which the structures are rasterized: the organs overwrite the skin and the bones.
 73PAINT_ORDER = ["skin", "bones"] + [name for name in ORGAN_IDS if name not in ("skin", "bones")]
 74
 75
 76def _find_dicom_series(path):
 77    """Find all DICOM series below the given folder and group them by patient and modality.
 78
 79    This works for the folder layout of the NBIA Data Retriever ('<Collection>/<Patient>/<Study>/<Series>/*.dcm')
 80    as well as for the layout of `util.download_source_tcia` ('<dst>/<SeriesInstanceUID>/*.dcm').
 81    """
 82    import pydicom
 83
 84    series = defaultdict(dict)
 85    for root, dirs, files in os.walk(path):
 86        dirs[:] = [d for d in dirs if d != "preprocessed"]
 87        dcm_files = [f for f in files if f.endswith(".dcm")]
 88        if not dcm_files:
 89            continue
 90        header = pydicom.dcmread(os.path.join(root, dcm_files[0]), specific_tags=["PatientID", "Modality"])
 91        series[str(header.PatientID)][str(header.Modality)] = root
 92    return series
 93
 94
 95def _load_dicom_volume(series_dir):
 96    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
 97
 98    Returns the volume in Hounsfield units and the geometry needed to map patient coordinates to pixel indices.
 99    """
100    import pydicom
101
102    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
103    slices.sort(key=lambda dcm: float(dcm.ImagePositionPatient[2]))
104
105    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
106    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
107    volume = np.round(volume).astype("int16")
108
109    geometry = {
110        "sop_uids": {str(dcm.SOPInstanceUID): i for i, dcm in enumerate(slices)},
111        "z_positions": {round(float(dcm.ImagePositionPatient[2]), 2): i for i, dcm in enumerate(slices)},
112        "origins": np.array([[float(v) for v in dcm.ImagePositionPatient] for dcm in slices]),
113        "orientation": np.array([float(v) for v in slices[0].ImageOrientationPatient]),
114        "spacing": np.array([float(v) for v in slices[0].PixelSpacing]),
115    }
116    return volume, geometry
117
118
119def _select_rtstruct(rtstruct_dir):
120    """Select the RTSTRUCT file of a series.
121
122    103 series contain two RTSTRUCT files, which only differ in their skin contours. We use the most recent one.
123    """
124    import pydicom
125
126    rtstruct_paths = glob(os.path.join(rtstruct_dir, "*.dcm"))
127    if len(rtstruct_paths) == 1:
128        return rtstruct_paths[0]
129    creation_dates = [
130        str(pydicom.dcmread(p, specific_tags=["InstanceCreationDate"]).InstanceCreationDate) for p in rtstruct_paths
131    ]
132    return rtstruct_paths[int(np.argmax(creation_dates))]
133
134
135def _rasterize_rtstruct(rtstruct_path, shape, geometry):
136    """Rasterize the closed planar contours of a DICOM RTSTRUCT on the grid of the reference CT volume.
137
138    Each contour is mapped to its CT slice via the referenced SOP instance (or its z position) and filled with
139    `cv2.fillPoly` (as in rt-utils). Multiple contours of a structure on the same slice are combined with XOR,
140    so that inner contours are treated as holes. The structures are then painted in the order given by `PAINT_ORDER`.
141    """
142    import cv2
143    import pydicom
144
145    rtstruct = pydicom.dcmread(rtstruct_path)
146    roi_names = {int(roi.ROINumber): str(roi.ROIName) for roi in rtstruct.StructureSetROISequence}
147
148    row_dir, col_dir = geometry["orientation"][3:], geometry["orientation"][:3]
149    row_spacing, col_spacing = geometry["spacing"]
150
151    masks = {}
152    for roi_contour in rtstruct.ROIContourSequence:
153        roi_name = roi_names[int(roi_contour.ReferencedROINumber)]
154        if roi_name not in ROI_NAMES or "ContourSequence" not in roi_contour:
155            continue
156        organ = ROI_NAMES[roi_name]
157        mask = masks.setdefault(organ, np.zeros(shape, dtype="bool"))
158
159        for contour in roi_contour.ContourSequence:
160            assert contour.ContourGeometricType == "CLOSED_PLANAR", f"Unexpected contour type in {rtstruct_path}."
161            points = np.array(contour.ContourData, dtype="float64").reshape(-1, 3)
162
163            if "ContourImageSequence" in contour:
164                z = geometry["sop_uids"].get(str(contour.ContourImageSequence[0].ReferencedSOPInstanceUID))
165            else:
166                z = geometry["z_positions"].get(round(float(points[0, 2]), 2))
167            if z is None:
168                warn(f"Skipping a contour of '{roi_name}' in {rtstruct_path}, which does not match a CT slice.")
169                continue
170
171            offsets = points - geometry["origins"][z]
172            rows = offsets @ row_dir / row_spacing
173            cols = offsets @ col_dir / col_spacing
174            contour_mask = np.zeros(shape[1:], dtype="uint8")
175            cv2.fillPoly(contour_mask, [np.round(np.stack([cols, rows], axis=1)).astype("int32")], 1)
176            mask[z] ^= contour_mask.astype("bool")
177
178    labels = np.zeros(shape, dtype="uint8")
179    for organ in PAINT_ORDER:
180        if organ in masks:
181            labels[masks[organ]] = ORGAN_IDS[organ]
182    return labels
183
184
185def _preprocess_pediatric_ct_seg(series, preprocessed_dir):
186    import h5py
187
188    os.makedirs(preprocessed_dir, exist_ok=True)
189    for patient_id, series_dirs in tqdm(sorted(series.items()), desc="Preprocess Pediatric-CT-SEG"):
190        out_path = os.path.join(preprocessed_dir, f"{patient_id}.h5")
191        if os.path.exists(out_path):
192            continue
193
194        volume, geometry = _load_dicom_volume(series_dirs["CT"])
195        rtstruct_path = _select_rtstruct(series_dirs["RTSTRUCT"])
196        labels = _rasterize_rtstruct(rtstruct_path, volume.shape, geometry)
197
198        with h5py.File(out_path, "w") as f:
199            f.create_dataset("raw", data=volume, compression="gzip")
200            f.create_dataset("labels", data=labels, compression="gzip")
201
202
203def get_pediatric_ct_seg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
204    """Download the Pediatric-CT-SEG dataset.
205
206    Args:
207        path: Filepath to a folder where the data is downloaded for further processing.
208        download: Whether to download the data if it is not present.
209
210    Returns:
211        Filepath where the preprocessed data is stored.
212    """
213    preprocessed_dir = os.path.join(path, "preprocessed")
214    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == NUM_VOLUMES:
215        return preprocessed_dir
216
217    os.makedirs(path, exist_ok=True)
218
219    # Download the DICOM series (CT and RTSTRUCT) from the TCIA manifest, unless they are present already
220    # (e.g. downloaded with the NBIA Data Retriever).
221    series = _find_dicom_series(path)
222    if not series:
223        util.download_source_tcia(
224            path=os.path.join(path, "Pediatric-CT-SEG-Mar-22-2022-manifest.tcia"), url=URL,
225            dst=os.path.join(path, "dicom"), csv_filename=os.path.join(path, "pediatric_ct_seg_series"),
226            download=download,
227        )
228        series = _find_dicom_series(path)
229
230    missing = [pid for pid, series_dirs in series.items() if {"CT", "RTSTRUCT"} - set(series_dirs)]
231    assert not missing, f"The CT or RTSTRUCT series is missing for the patients {missing}."
232
233    _preprocess_pediatric_ct_seg(series, preprocessed_dir)
234    return preprocessed_dir
235
236
237def get_pediatric_ct_seg_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
238    """Get paths to the Pediatric-CT-SEG data.
239
240    Args:
241        path: Filepath to a folder where the data is downloaded for further processing.
242        download: Whether to download the data if it is not present.
243
244    Returns:
245        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
246    """
247    data_dir = get_pediatric_ct_seg_data(path, download)
248    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
249    return volume_paths
250
251
252def get_pediatric_ct_seg_dataset(
253    path: Union[os.PathLike, str],
254    patch_shape: Tuple[int, ...],
255    resize_inputs: bool = False,
256    download: bool = False,
257    **kwargs
258) -> Dataset:
259    """Get the Pediatric-CT-SEG dataset for organ segmentation.
260
261    Args:
262        path: Filepath to a folder where the data is downloaded for further processing.
263        patch_shape: The patch shape to use for training.
264        resize_inputs: Whether to resize inputs to the desired patch shape.
265        download: Whether to download the data if it is not present.
266        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
267
268    Returns:
269        The segmentation dataset.
270    """
271    volume_paths = get_pediatric_ct_seg_paths(path, download)
272
273    if resize_inputs:
274        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
275        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
276            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
277        )
278
279    return torch_em.default_segmentation_dataset(
280        raw_paths=volume_paths,
281        raw_key="raw",
282        label_paths=volume_paths,
283        label_key="labels",
284        patch_shape=patch_shape,
285        is_seg_dataset=True,
286        **kwargs
287    )
288
289
290def get_pediatric_ct_seg_loader(
291    path: Union[os.PathLike, str],
292    batch_size: int,
293    patch_shape: Tuple[int, ...],
294    resize_inputs: bool = False,
295    download: bool = False,
296    **kwargs
297) -> DataLoader:
298    """Get the Pediatric-CT-SEG dataloader for organ segmentation.
299
300    Args:
301        path: Filepath to a folder where the data is downloaded for further processing.
302        batch_size: The batch size for training.
303        patch_shape: The patch shape to use for training.
304        resize_inputs: Whether to resize inputs to the desired patch shape.
305        download: Whether to download the data if it is not present.
306        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
307
308    Returns:
309        The DataLoader.
310    """
311    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
312    dataset = get_pediatric_ct_seg_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
313    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)