torch_em.data.datasets.medical.vs_mc_rc

  1"""The Vestibular-Schwannoma-MC-RC dataset contains annotations for the vestibular schwannoma tumor in MRI.
  2
  3It consists of routine clinical T1-weighted and T2-weighted MRI series of patients with a single sporadic
  4vestibular schwannoma, acquired at 10 sites in the United Kingdom. The tumor is annotated as a binary mask
  5(label id 1) in 301 series of 124 patients, each series belongs to one timepoint of a patient. Version 2 of the
  6collection also provides the images of 126 further timepoints, which belong to the validation and test set of a
  7challenge and have no public annotations, so they are not used here.
  8
  9The collection has no official data splits. A split is derived from the patient ids (sorted, first 80% train,
 10next 10% val, last 10% test), so that all timepoints of a patient are in the same split.
 11
 12The dataset is located at https://doi.org/10.7937/HRZH-2N82 and is distributed under the CC BY 4.0 license.
 13This dataset is from the publication https://doi.org/10.3389/fncom.2024.1365727.
 14Please cite it if you use this dataset in your research.
 15"""
 16
 17import os
 18from glob import glob
 19from tqdm import tqdm
 20from natsort import natsorted
 21from typing import Union, Tuple, List, Literal
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = (
 31    "https://www.cancerimagingarchive.net/wp-content/uploads/"
 32    "VS-MC-RC-Segmentations_with_pathname_spreadsheet-Release2023.zip"
 33)
 34CHECKSUM = "6cbbbd10552398795bdac868385aceac70f47817bcf012d126424831ecaf198e"
 35
 36
 37def _get_labeled_series(segmentation_dir):
 38    import pandas as pd
 39
 40    info_path = glob(os.path.join(segmentation_dir, "*", "*segmentation-info*.csv"))[0]
 41    label_dir = glob(os.path.join(segmentation_dir, "*", "VS-MC-RC segmentations*"))[0]
 42    series = pd.read_csv(info_path)
 43    series = series[series["SegmentationPath"].notna()]
 44    # The segmentations of the timepoints of the challenge validation and test set are not public.
 45    series = series[[os.path.exists(os.path.join(label_dir, p)) for p in series["SegmentationPath"]]]
 46    return series, label_dir
 47
 48
 49def _resample_labels(nifti, geometry, shape):
 50    """Resample the labels onto the voxel grid of the DICOM series, as their axes can be ordered differently."""
 51    import numpy as np
 52    from scipy.ndimage import map_coordinates
 53
 54    labels = (np.asarray(nifti.dataobj) > 0).astype("float32")
 55    ras_to_voxel = np.linalg.inv(nifti.affine)
 56    row_step = geometry["spacing"][1] * geometry["row_direction"]
 57    column_step = geometry["spacing"][0] * geometry["column_direction"]
 58
 59    ys, xs = np.meshgrid(np.arange(shape[1]), np.arange(shape[2]), indexing="ij")
 60    in_plane = ys[..., None] * column_step + xs[..., None] * row_step  # (y, x, 3)
 61    resampled = np.zeros(shape, dtype="uint8")
 62    for z, origin in enumerate(geometry["origin"]):
 63        lps = in_plane + origin
 64        ras = lps * np.array([-1.0, -1.0, 1.0])
 65        voxel = ras @ ras_to_voxel[:3, :3].T + ras_to_voxel[:3, 3]
 66        resampled[z] = map_coordinates(labels, voxel.transpose(2, 0, 1), order=0, mode="constant") > 0.5
 67    return resampled
 68
 69
 70def _preprocess_vs_mc_rc(series, label_dir, dicom_dir, preprocessed_dir):
 71    import h5py
 72    import nibabel as nib
 73
 74    os.makedirs(preprocessed_dir, exist_ok=True)
 75    for _, row in tqdm(series.iterrows(), total=len(series), desc="Preprocess Vestibular-Schwannoma-MC-RC"):
 76        patient_id, date, label_name = row["SegmentationPath"].split("/")
 77        modality = label_name.split("_")[-1].split(".")[0]
 78        out_path = os.path.join(preprocessed_dir, f"{patient_id}_{date}_{modality}.h5")
 79        if os.path.exists(out_path):
 80            continue
 81
 82        volume, geometry = util.load_dicom_series(os.path.join(dicom_dir, row["series_instance_uid"]))
 83        labels = _resample_labels(nib.load(os.path.join(label_dir, row["SegmentationPath"])), geometry, volume.shape)
 84
 85        # The file is written to a temporary path first, so an interrupted run is not mistaken for a finished one.
 86        tmp_path = f"{out_path}.incomplete"
 87        with h5py.File(tmp_path, "w") as f:
 88            f.create_dataset("raw", data=volume, compression="gzip")
 89            f.create_dataset("labels", data=labels, compression="gzip")
 90        os.replace(tmp_path, out_path)
 91
 92
 93def get_vs_mc_rc_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 94    """Download the Vestibular-Schwannoma-MC-RC dataset.
 95
 96    NOTE: This requires the pydicom and nibabel python packages. The DICOM series of the annotated timepoints are
 97    about 9 GB.
 98
 99    Args:
100        path: Filepath to a folder where the data is downloaded for further processing.
101        download: Whether to download the data if it is not present.
102
103    Returns:
104        Filepath where the preprocessed data is stored.
105    """
106    os.makedirs(path, exist_ok=True)
107
108    segmentation_dir = os.path.join(path, "segmentations")
109    if not os.path.exists(segmentation_dir):
110        zip_path = os.path.join(path, "segmentations.zip")
111        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
112        util.unzip(zip_path=zip_path, dst=segmentation_dir, remove=False)
113
114    series, label_dir = _get_labeled_series(segmentation_dir)
115
116    dicom_dir = os.path.join(path, "dicom")
117    uids = sorted(series["series_instance_uid"].unique())
118    if download:
119        util.download_tcia_series(uids, dst=dicom_dir, csv_filename=os.path.join(path, "vs_mc_rc_series"))
120    elif not all(glob(os.path.join(dicom_dir, uid, "*.dcm")) for uid in uids):
121        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
122
123    preprocessed_dir = os.path.join(path, "preprocessed")
124    _preprocess_vs_mc_rc(series, label_dir, dicom_dir, preprocessed_dir)
125    return preprocessed_dir
126
127
128def get_vs_mc_rc_paths(
129    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False,
130) -> List[str]:
131    """Get paths to the Vestibular-Schwannoma-MC-RC data.
132
133    Args:
134        path: Filepath to a folder where the data is downloaded for further processing.
135        split: The data split to use. Either 'train', 'val' or 'test'.
136        download: Whether to download the data if it is not present.
137
138    Returns:
139        List of filepaths for the stored data.
140    """
141    assert split in ("train", "val", "test"), f"'{split}' is not a valid data split."
142
143    preprocessed_dir = get_vs_mc_rc_data(path, download)
144    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
145    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
146
147    patient_ids = sorted({os.path.basename(p).split("_")[0] for p in volume_paths})
148    n_train, n_val = int(0.8 * len(patient_ids)), int(0.1 * len(patient_ids))
149    split_ids = {
150        "train": patient_ids[:n_train],
151        "val": patient_ids[n_train:n_train + n_val],
152        "test": patient_ids[n_train + n_val:],
153    }[split]
154    return [p for p in volume_paths if os.path.basename(p).split("_")[0] in split_ids]
155
156
157def get_vs_mc_rc_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, ...],
160    split: Literal["train", "val", "test"],
161    resize_inputs: bool = False,
162    download: bool = False,
163    **kwargs
164) -> Dataset:
165    """Get the Vestibular-Schwannoma-MC-RC dataset for tumor segmentation.
166
167    Args:
168        path: Filepath to a folder where the data is downloaded for further processing.
169        patch_shape: The patch shape to use for training.
170        split: The data split to use. Either 'train', 'val' or 'test'.
171        resize_inputs: Whether to resize inputs to the desired patch shape.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
174
175    Returns:
176        The segmentation dataset.
177    """
178    volume_paths = get_vs_mc_rc_paths(path, split, download)
179
180    if resize_inputs:
181        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
182        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
183            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
184        )
185
186    return torch_em.default_segmentation_dataset(
187        raw_paths=volume_paths,
188        raw_key="raw",
189        label_paths=volume_paths,
190        label_key="labels",
191        patch_shape=patch_shape,
192        is_seg_dataset=True,
193        **kwargs
194    )
195
196
197def get_vs_mc_rc_loader(
198    path: Union[os.PathLike, str],
199    batch_size: int,
200    patch_shape: Tuple[int, ...],
201    split: Literal["train", "val", "test"],
202    resize_inputs: bool = False,
203    download: bool = False,
204    **kwargs
205) -> DataLoader:
206    """Get the Vestibular-Schwannoma-MC-RC dataloader for tumor segmentation.
207
208    Args:
209        path: Filepath to a folder where the data is downloaded for further processing.
210        batch_size: The batch size for training.
211        patch_shape: The patch shape to use for training.
212        split: The data split to use. Either 'train', 'val' or 'test'.
213        resize_inputs: Whether to resize inputs to the desired patch shape.
214        download: Whether to download the data if it is not present.
215        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
216
217    Returns:
218        The DataLoader.
219    """
220    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
221    dataset = get_vs_mc_rc_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
222    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)