torch_em.data.datasets.medical.ctspinopelvic1k

  1"""The CTSpinoPelvic1K dataset contains annotations for the spine, pelvis, ribs and femora in CT scans.
  2
  3The dataset consists of 802 CT COLONOGRAPHY scans from TCIA with 69-class annotations (see
  4`CLASS_NAMES`) in one coordinate frame: the cervical to lumbar vertebrae (with a sixth lumbar
  5vertebra where present), the sacrum, coccyx, hips, femora, the individual ribs (per side and per
  6level, including a rib on a lumbar vertebra where present) and spinal hardware. The vertebral
  7annotations derive from CTSpine1K and the pelvic annotations from CTPelvic1K, remapped onto one
  8coordinate frame and paired with the exact CT series they were drawn on, which neither of those
  9releases published. The manifest also carries a lumbosacral transitional anatomy label and a
 10Castellvi grade per case.
 11
 12NOTE: The images are not distributed with the release: they are 193 GB against 1.8 GB of labels
 13and already public on TCIA, so this module downloads and reconstructs them, following the same
 14steps and settings as the release. A raw DICOM conversion (e.g. with dcm2niix) does not share the
 15label's grid, so it is resampled onto it: DICOM patient coordinates are LPS, the label is stored
 16in a RAS+ world frame (as dcm2niix would produce), and a trilinear resampling with a -1024 HU fill
 17value outside the original extent completes the match.
 18
 19NOTE: This requires the pydicom and scipy python packages.
 20
 21The dataset is located at https://doi.org/10.5281/zenodo.22642578 and is distributed under the
 22CC BY-NC-SA 4.0 license for research use only.
 23This dataset is from the publications https://doi.org/10.48550/arXiv.2105.14711 (CTSpine1K) and
 24https://doi.org/10.1007/s11548-021-02363-8 (CTPelvic1K); please cite the Zenodo record and both
 25publications if you use this dataset in your research.
 26"""
 27
 28import os
 29import json
 30from glob import glob
 31from tqdm import tqdm
 32from natsort import natsorted
 33from typing import Union, Tuple, List
 34
 35import numpy as np
 36
 37from torch.utils.data import Dataset, DataLoader
 38
 39import torch_em
 40
 41from .adrenal_acc import _load_dicom_volume
 42from .. import util
 43
 44
 45URLS = {
 46    "manifest": "https://zenodo.org/records/22642578/files/manifest.json?download=1",
 47    "labels": "https://zenodo.org/records/22642578/files/labels.zip?download=1",
 48}
 49
 50CHECKSUMS = {
 51    "manifest": "6f25aac0ea6f4b46801d372f3d4762f05cd0dceb05da04ac73a3ad54dd427c55",
 52    "labels": "a6d0df210fea4660095dc27d9caf28120ceb35313daeb2b5762587fd458383f6",
 53}
 54
 55IGNORE_LABEL = 255
 56"""The sentinel that marks a voxel excluded from the annotation, e.g. outside a partial scan."""
 57
 58CLASS_NAMES = [
 59    "C1", "C2", "C3", "C4", "C5", "C6", "C7",
 60    "T1", "T2", "T3", "T4", "T5", "T6", "T7", "T8", "T9", "T10", "T11", "T12",
 61    "L1", "L2", "L3", "L4", "L5", "L6",
 62    "sacrum", "coccyx", "T13", "S1", "left_hip", "right_hip", "femur_left", "femur_right",
 63] + [f"rib_left_{i}" for i in range(1, 14)] + [f"rib_right_{i}" for i in range(1, 14)] + [
 64    "rib_left_lumbar", "rib_right_lumbar", "hardware", "hardware_cage", "hardware_screw_rod",
 65    "hardware_plate", "hardware_arthroplasty", "hardware_si_screw", "hardware_osteosynthesis",
 66]
 67"""The 68 foreground classes of the CTSpinoPelvic1K dataset. The label id of a class is its
 681-based index; 255 marks an excluded voxel rather than a class. See also `CLASS_IDS`."""
 69
 70CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 71"""Mapping from the anatomical structure to its label id."""
 72
 73
 74def _series_uid(record):
 75    """The UID the labels of a record were drawn on: the spine series, or the pelvic one for the
 76    pelvis-only records that have no spine annotation.
 77
 78    A handful of records suffix the UID with '_orientation_fixed', a note from the release's own
 79    pipeline rather than part of the UID itself, which is stripped to get a downloadable series UID.
 80    """
 81    uid = str(record.get("spine_series_uid") or "").strip() or str(record.get("pelvic_series_uid") or "").strip()
 82    return uid.removesuffix("_orientation_fixed")
 83
 84
 85def _resample_to_label(volume, ct_affine, label_shape, label_affine):
 86    """Resample a DICOM-derived volume onto the grid of its label, matching the release's own
 87    reconstruction: DICOM patient coordinates are LPS, converted to the RAS+ frame of the label by
 88    negating x and y, then a trilinear resampling with a -1024 HU fill value outside the CT extent.
 89    """
 90    from scipy.ndimage import affine_transform
 91
 92    lps_to_ras = np.diag([-1.0, -1.0, 1.0, 1.0])
 93    ras_affine = lps_to_ras @ ct_affine
 94    if volume.shape == label_shape and np.allclose(ras_affine, label_affine, atol=1e-3):
 95        return volume.astype("int16")
 96
 97    to_ct_index = np.linalg.inv(ras_affine) @ label_affine
 98    resampled = affine_transform(
 99        volume.astype("float32"), to_ct_index[:3, :3], offset=to_ct_index[:3, 3],
100        output_shape=label_shape, order=1, mode="constant", cval=-1024.0,
101    )
102    return np.round(resampled).astype("int16")
103
104
105def _preprocess_ctspinopelvic1k(label_dir, manifest_path, dicom_dir, preprocessed_dir):
106    import h5py
107    import nibabel as nib
108
109    with open(manifest_path) as f:
110        manifest = json.load(f)
111    records = manifest if isinstance(manifest, list) else manifest.get("records", list(manifest.values()))
112
113    os.makedirs(preprocessed_dir, exist_ok=True)
114    for record in tqdm(records, desc="Preprocess CTSpinoPelvic1K"):
115        case_id = os.path.basename(record["label_file"]).split("_")[0]
116        out_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
117        if os.path.exists(out_path):
118            continue
119
120        label_path = os.path.join(label_dir, f"{case_id}_label.nii.gz")
121        if not os.path.exists(label_path):
122            continue
123
124        series_dir = os.path.join(dicom_dir, _series_uid(record))
125        if not glob(os.path.join(series_dir, "*.dcm")):
126            continue
127
128        label_image = nib.load(label_path)
129        labels = np.asarray(label_image.dataobj)
130        volume, ct_affine = _load_dicom_volume(series_dir)
131        raw = _resample_to_label(volume, ct_affine, labels.shape, label_image.affine)
132
133        with h5py.File(out_path, "w") as f:
134            f.create_dataset("raw", data=raw, compression="gzip")
135            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
136
137
138def get_ctspinopelvic1k_data(path: Union[os.PathLike, str], download: bool = False) -> str:
139    """Download the CTSpinoPelvic1K dataset.
140
141    The images are reconstructed from TCIA, which is several hundred gigabytes and can take many
142    hours to download depending on the connection to the TCIA servers.
143
144    Args:
145        path: Filepath to a folder where the data is downloaded for further processing.
146        download: Whether to download the data if it is not present.
147
148    Returns:
149        Filepath where the preprocessed data is stored.
150    """
151    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
152    preprocessed_dir = os.path.join(path, "preprocessed")
153
154    os.makedirs(path, exist_ok=True)
155
156    manifest_path = os.path.join(path, "manifest.json")
157    util.download_source(
158        path=manifest_path, url=URLS["manifest"], download=download, checksum=CHECKSUMS["manifest"]
159    )
160
161    label_dir = os.path.join(path, "labels")
162    if not os.path.exists(label_dir):
163        zip_path = os.path.join(path, "labels.zip")
164        util.download_source(path=zip_path, url=URLS["labels"], download=download, checksum=CHECKSUMS["labels"])
165        util.unzip(zip_path=zip_path, dst=path, remove=False)
166
167    with open(manifest_path) as f:
168        manifest = json.load(f)
169    records = manifest if isinstance(manifest, list) else manifest.get("records", list(manifest.values()))
170    series_uids = sorted({_series_uid(record) for record in records if _series_uid(record)})
171
172    dicom_dir = os.path.join(path, "dicom")
173    if download:  # Series that were downloaded already are skipped.
174        util.download_tcia_series(series_uids, dst=dicom_dir, csv_filename=os.path.join(path, "ctspinopelvic1k"))
175    elif not all(glob(os.path.join(dicom_dir, uid, "*.dcm")) for uid in series_uids):
176        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
177
178    _preprocess_ctspinopelvic1k(label_dir, manifest_path, dicom_dir, preprocessed_dir)
179    return preprocessed_dir
180
181
182def get_ctspinopelvic1k_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
183    """Get paths to the CTSpinoPelvic1K data.
184
185    Args:
186        path: Filepath to a folder where the data is downloaded for further processing.
187        download: Whether to download the data if it is not present.
188
189    Returns:
190        List of filepaths for the stored data.
191    """
192    preprocessed_dir = get_ctspinopelvic1k_data(path, download)
193    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
194    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
195    return volume_paths
196
197
198def get_ctspinopelvic1k_dataset(
199    path: Union[os.PathLike, str],
200    patch_shape: Tuple[int, ...],
201    resize_inputs: bool = False,
202    download: bool = False,
203    **kwargs
204) -> Dataset:
205    """Get the CTSpinoPelvic1K dataset for spine, pelvis, rib and femur segmentation.
206
207    Args:
208        path: Filepath to a folder where the data is downloaded for further processing.
209        patch_shape: The patch shape to use for training.
210        resize_inputs: Whether to resize inputs to the desired patch shape.
211        download: Whether to download the data if it is not present.
212        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
213
214    Returns:
215        The segmentation dataset.
216    """
217    volume_paths = get_ctspinopelvic1k_paths(path, download)
218
219    if resize_inputs:
220        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
221        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
222            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
223        )
224
225    return torch_em.default_segmentation_dataset(
226        raw_paths=volume_paths,
227        raw_key="raw",
228        label_paths=volume_paths,
229        label_key="labels",
230        patch_shape=patch_shape,
231        is_seg_dataset=True,
232        **kwargs
233    )
234
235
236def get_ctspinopelvic1k_loader(
237    path: Union[os.PathLike, str],
238    batch_size: int,
239    patch_shape: Tuple[int, ...],
240    resize_inputs: bool = False,
241    download: bool = False,
242    **kwargs
243) -> DataLoader:
244    """Get the CTSpinoPelvic1K dataloader for spine, pelvis, rib and femur segmentation.
245
246    Args:
247        path: Filepath to a folder where the data is downloaded for further processing.
248        batch_size: The batch size for training.
249        patch_shape: The patch shape to use for training.
250        resize_inputs: Whether to resize inputs to the desired patch shape.
251        download: Whether to download the data if it is not present.
252        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
253
254    Returns:
255        The DataLoader.
256    """
257    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
258    dataset = get_ctspinopelvic1k_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
259    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)