torch_em.data.datasets.medical.spine_mets

  1"""The Spine-Mets-CT-SEG dataset contains annotations for vertebra segmentation in CT of patients
  2with spinal metastases (pre and post radiotherapy).
  3
  4It consists of 55 CT volumes with semantic labels for the individual vertebrae. The CT scans are distributed
  5as DICOM series and the labels as DICOM-SEG objects, which are converted and stored in hdf5 files by this module.
  6The vertebra labels are mapped to consistent semantic ids: 1-7: C1-C7, 8-19: T1-T12, 20-24: L1-L5, 25: S1.
  7
  8NOTE: This requires the pydicom python package.
  9
 10The dataset is located at https://www.cancerimagingarchive.net/collection/spine-mets-ct-seg/.
 11
 12This dataset is from the publication https://doi.org/10.7937/kh36-ds04.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17import csv
 18from glob import glob
 19from tqdm import tqdm
 20from natsort import natsorted
 21from collections import defaultdict
 22from typing import Union, Tuple, List
 23
 24import numpy as np
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL = "https://www.cancerimagingarchive.net/wp-content/uploads/Spine-Mets-CT-SEG_v1_2024.tcia"
 34
 35# The DICOM series are downloaded individually from TCIA.
 36CHECKSUM = None
 37
 38VERTEBRA_IDS = {
 39    **{f"C{i}": i for i in range(1, 8)},
 40    **{f"T{i}": 7 + i for i in range(1, 13)},
 41    **{f"L{i}": 19 + i for i in range(1, 6)},
 42    "S1": 25,
 43}
 44
 45
 46def _load_dicom_volume(series_dir):
 47    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
 48
 49    Returns the volume in Hounsfield units, the z position of each slice and the image orientation
 50    (DICOM 'ImageOrientationPatient').
 51    """
 52    import pydicom
 53
 54    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 55    slices.sort(key=lambda dcm: float(dcm.ImagePositionPatient[2]))
 56
 57    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 58    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 59    volume = np.round(volume).astype("int16")
 60
 61    z_positions = np.round([float(dcm.ImagePositionPatient[2]) for dcm in slices], 3)
 62    orientation = np.round([float(v) for v in slices[0].ImageOrientationPatient]).astype("int")
 63    return volume, z_positions, orientation
 64
 65
 66def _load_dicom_seg(seg_path, shape, z_positions, orientation):
 67    """Convert a DICOM-SEG object into a semantic label volume aligned with the reference CT volume."""
 68    import pydicom
 69
 70    seg = pydicom.dcmread(seg_path)
 71    frames = seg.pixel_array
 72    if frames.ndim == 2:  # A segmentation with a single frame.
 73        frames = frames[None]
 74
 75    segment_ids = {}
 76    for segment in seg.SegmentSequence:
 77        name = str(segment.SegmentLabel).replace("vertebra", "").strip()
 78        if name not in VERTEBRA_IDS:
 79            raise ValueError(f"Unknown segment label '{segment.SegmentLabel}' in {seg_path}.")
 80        segment_ids[int(segment.SegmentNumber)] = VERTEBRA_IDS[name]
 81
 82    # The segmentation frames may use a different in-plane orientation than the CT slices,
 83    # in which case they have to be flipped to align them.
 84    seg_orientation = seg.SharedFunctionalGroupsSequence[0].PlaneOrientationSequence[0].ImageOrientationPatient
 85    seg_orientation = np.round([float(v) for v in seg_orientation]).astype("int")
 86    assert np.all(np.abs(seg_orientation) == np.abs(orientation)), f"Unexpected orientation in {seg_path}."
 87    flip_axes = []
 88    if np.any(seg_orientation[3:] != orientation[3:]):  # The direction of the rows differs.
 89        flip_axes.append(0)
 90    if np.any(seg_orientation[:3] != orientation[:3]):  # The direction of the columns differs.
 91        flip_axes.append(1)
 92
 93    slice_ids = {z: i for i, z in enumerate(z_positions)}
 94    labels = np.zeros(shape, dtype="uint8")
 95    for frame, frame_group in zip(frames, seg.PerFrameFunctionalGroupsSequence):
 96        z = round(float(frame_group.PlanePositionSequence[0].ImagePositionPatient[2]), 3)
 97        segment_number = int(frame_group.SegmentIdentificationSequence[0].ReferencedSegmentNumber)
 98        mask = frame.astype("bool")
 99        if flip_axes:
100            mask = np.flip(mask, axis=flip_axes)
101        labels[slice_ids[z]][mask] = segment_ids[segment_number]
102
103    return labels
104
105
106def _preprocess_spine_mets(dicom_dir, csv_path, preprocessed_dir):
107    import h5py
108
109    series_per_subject = defaultdict(dict)
110    with open(csv_path, "r") as f:
111        for row in csv.DictReader(f):
112            series_per_subject[row["Subject ID"]][row["Modality"]] = os.path.join(dicom_dir, row["Series UID"])
113
114    os.makedirs(preprocessed_dir, exist_ok=True)
115    for subject_id, series_dirs in tqdm(sorted(series_per_subject.items()), desc="Preprocess Spine-Mets-CT-SEG"):
116        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
117        if os.path.exists(out_path):
118            continue
119
120        volume, z_positions, orientation = _load_dicom_volume(series_dirs["CT"])
121        seg_path = glob(os.path.join(series_dirs["SEG"], "*.dcm"))[0]
122        labels = _load_dicom_seg(seg_path, volume.shape, z_positions, orientation)
123
124        with h5py.File(out_path, "w") as f:
125            f.create_dataset("raw", data=volume, compression="gzip")
126            f.create_dataset("labels", data=labels, compression="gzip")
127
128
129def get_spine_mets_data(path: Union[os.PathLike, str], download: bool = False) -> str:
130    """Download the Spine-Mets-CT-SEG dataset.
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        Filepath where the preprocessed data is stored.
138    """
139    preprocessed_dir = os.path.join(path, "preprocessed")
140    if os.path.exists(preprocessed_dir):
141        return preprocessed_dir
142
143    os.makedirs(path, exist_ok=True)
144
145    # Download the DICOM series (CT and SEG) from the TCIA manifest.
146    dicom_dir = os.path.join(path, "dicom")
147    csv_path = os.path.join(path, "spine_mets_series")
148    util.download_source_tcia(
149        path=os.path.join(path, "Spine-Mets-CT-SEG_v1_2024.tcia"), url=URL, dst=dicom_dir,
150        csv_filename=csv_path, download=download,
151    )
152
153    _preprocess_spine_mets(dicom_dir, f"{csv_path}.csv", preprocessed_dir)
154    return preprocessed_dir
155
156
157def get_spine_mets_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
158    """Get paths to the Spine-Mets-CT-SEG data.
159
160    Args:
161        path: Filepath to a folder where the data is downloaded for further processing.
162        download: Whether to download the data if it is not present.
163
164    Returns:
165        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
166    """
167    data_dir = get_spine_mets_data(path, download)
168    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
169    return volume_paths
170
171
172def get_spine_mets_dataset(
173    path: Union[os.PathLike, str],
174    patch_shape: Tuple[int, ...],
175    resize_inputs: bool = False,
176    download: bool = False,
177    **kwargs
178) -> Dataset:
179    """Get the Spine-Mets-CT-SEG dataset for vertebra segmentation.
180
181    Args:
182        path: Filepath to a folder where the data is downloaded for further processing.
183        patch_shape: The patch shape to use for training.
184        resize_inputs: Whether to resize inputs to the desired patch shape.
185        download: Whether to download the data if it is not present.
186        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
187
188    Returns:
189        The segmentation dataset.
190    """
191    volume_paths = get_spine_mets_paths(path, download)
192
193    if resize_inputs:
194        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
195        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
196            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
197        )
198
199    return torch_em.default_segmentation_dataset(
200        raw_paths=volume_paths,
201        raw_key="raw",
202        label_paths=volume_paths,
203        label_key="labels",
204        patch_shape=patch_shape,
205        is_seg_dataset=True,
206        **kwargs
207    )
208
209
210def get_spine_mets_loader(
211    path: Union[os.PathLike, str],
212    batch_size: int,
213    patch_shape: Tuple[int, ...],
214    resize_inputs: bool = False,
215    download: bool = False,
216    **kwargs
217) -> DataLoader:
218    """Get the Spine-Mets-CT-SEG dataloader for vertebra segmentation.
219
220    Args:
221        path: Filepath to a folder where the data is downloaded for further processing.
222        batch_size: The batch size for training.
223        patch_shape: The patch shape to use for training.
224        resize_inputs: Whether to resize inputs to the desired patch shape.
225        download: Whether to download the data if it is not present.
226        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
227
228    Returns:
229        The DataLoader.
230    """
231    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
232    dataset = get_spine_mets_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
233    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)