torch_em.data.datasets.medical.remind

  1"""The ReMIND dataset contains annotations for brain tumor and resection segmentation in pre- and
  2intra-operative MRI of patients who underwent image-guided tumor resection.
  3
  4The full collection consists of 114 patients with 369 pre-operative MRI, 301 intra-operative MRI and 320 3D
  5intra-operative ultrasound series. Segmentations are only provided for 221 MRI series (the ultrasound series are not
  6annotated), so this module downloads and converts only these MRI series and their segmentations (ca. 8.5 GB instead
  7of the full 43.5 GB collection). The MRI is distributed as DICOM series and the segmentations as DICOM-SEG objects,
  8which are converted and stored in hdf5 files by this module. The segmentations of all structures annotated for a MRI
  9series are merged into a single label volume with the semantic ids 1: cerebrum, 2: ventricles,
 103: previous resection cavity, 4: tumor, 5: tumor target, 6: residual tumor (structures with a higher id overwrite
 11structures with a lower id where they overlap). Only the structures deemed necessary for the surgery are segmented,
 12so most volumes contain only a subset of these structures.
 13
 14The MRI sequences are exposed via the 'modality' argument: 't1c' (contrast-enhanced T1), 't1' (native T1,
 15including MP2RAGE), 't2' and 'flair'. The 'study' argument selects pre-operative or intra-operative MRI.
 16
 17NOTE: This requires the pydicom python package.
 18
 19The dataset is located at https://www.cancerimagingarchive.net/collection/remind/.
 20
 21This dataset is from the publication https://doi.org/10.1038/s41597-024-03295-z.
 22The data was released at https://doi.org/10.7937/3RAG-D070.
 23Please cite it if you use this dataset in your research.
 24"""
 25
 26import os
 27import csv
 28import tempfile
 29from glob import glob
 30from tqdm import tqdm
 31from shutil import copyfileobj
 32from natsort import natsorted
 33from typing import Union, Tuple, List, Optional, Literal
 34
 35import numpy as np
 36import requests
 37
 38from torch.utils.data import Dataset, DataLoader
 39
 40import torch_em
 41
 42from .. import util
 43
 44
 45URL = "https://www.cancerimagingarchive.net/wp-content/uploads/ReMIND-Manifest-Sept-2023.tcia"
 46
 47# The DICOM series are downloaded individually from TCIA.
 48CHECKSUM = None
 49
 50LABEL_IDS = {
 51    "cerebrum": 1,
 52    "ventricles": 2,
 53    "previous_resection_cavity": 3,
 54    "tumor": 4,
 55    "tumor_target": 5,
 56    "tumor_residual": 6,
 57}
 58
 59MODALITIES = ["t1c", "t1", "t2", "flair"]
 60
 61
 62def _get_modality(series_description):
 63    """Derive the MRI sequence from the series description, e.g. '3D_AX_T1_postcontrast' or '2D_AX_T2_FLAIR'."""
 64    if "postcontrast" in series_description:
 65        return "t1c"
 66    elif "T1" in series_description:
 67        return "t1"
 68    elif "FLAIR" in series_description:
 69        return "flair"
 70    elif "T2" in series_description:
 71        return "t2"
 72    raise ValueError(f"Could not derive the modality from the series description '{series_description}'.")
 73
 74
 75def _load_dicom_volume(series_dir):
 76    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted along the slice normal.
 77
 78    Returns the volume and the affine matrix that maps voxel indices (z, y, x) to DICOM patient coordinates.
 79    """
 80    import pydicom
 81
 82    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 83    orientation = np.array([float(v) for v in slices[0].ImageOrientationPatient])
 84    row_dir, col_dir = orientation[:3], orientation[3:]
 85    normal = np.cross(row_dir, col_dir)
 86    slices.sort(key=lambda dcm: np.dot([float(v) for v in dcm.ImagePositionPatient], normal))
 87
 88    volume = np.stack([dcm.pixel_array for dcm in slices])
 89    if "RescaleSlope" in slices[0]:
 90        volume = volume.astype("float32") * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 91        volume = np.round(volume).astype("int16")
 92
 93    positions = np.array([[float(v) for v in dcm.ImagePositionPatient] for dcm in slices])
 94    spacing = [float(v) for v in slices[0].PixelSpacing]  # The spacing between rows and between columns.
 95    if len(slices) > 1:
 96        slice_step = (positions[-1] - positions[0]) / (len(slices) - 1)
 97    else:
 98        slice_step = normal * float(slices[0].SliceThickness)
 99
100    affine = np.eye(4)
101    affine[:3, 0] = slice_step
102    affine[:3, 1] = col_dir * spacing[0]
103    affine[:3, 2] = row_dir * spacing[1]
104    affine[:3, 3] = positions[0]
105    return volume, affine
106
107
108def _load_dicom_seg(seg_path):
109    """Load a DICOM-SEG object as a label volume with axes (z, y, x), where the segment number is used as label id.
110
111    Returns the label volume and the affine matrix that maps its voxel indices to DICOM patient coordinates.
112    """
113    import pydicom
114
115    seg = pydicom.dcmread(seg_path)
116    frames = seg.pixel_array
117    if frames.ndim == 2:  # A segmentation with a single frame.
118        frames = frames[None]
119
120    shared_group = seg.SharedFunctionalGroupsSequence[0]
121    orientation = np.array([float(v) for v in shared_group.PlaneOrientationSequence[0].ImageOrientationPatient])
122    row_dir, col_dir = orientation[:3], orientation[3:]
123    normal = np.cross(row_dir, col_dir)
124    pixel_measures = shared_group.PixelMeasuresSequence[0]
125    spacing = [float(v) for v in pixel_measures.PixelSpacing]  # The spacing between rows and between columns.
126
127    # The frames may be stored in arbitrary order and frames without any foreground may be skipped,
128    # so the position of each frame along the slice normal is derived from its patient position.
129    frame_groups = seg.PerFrameFunctionalGroupsSequence
130    positions = np.array([[float(v) for v in g.PlanePositionSequence[0].ImagePositionPatient] for g in frame_groups])
131    projections = positions @ normal
132    if "SpacingBetweenSlices" in pixel_measures:
133        slice_spacing = float(pixel_measures.SpacingBetweenSlices)
134    elif len(projections) > 1:
135        slice_spacing = np.min(np.diff(np.unique(np.round(projections, 3))))
136    else:
137        slice_spacing = float(pixel_measures.SliceThickness)
138    slice_ids = np.round((projections - projections.min()) / slice_spacing).astype("int")
139
140    labels = np.zeros((slice_ids.max() + 1, seg.Rows, seg.Columns), dtype="uint8")
141    for frame, frame_group, slice_id in zip(frames, frame_groups, slice_ids):
142        segment_number = int(frame_group.SegmentIdentificationSequence[0].ReferencedSegmentNumber)
143        labels[slice_id][frame.astype("bool")] = segment_number
144
145    affine = np.eye(4)
146    affine[:3, 0] = normal * slice_spacing
147    affine[:3, 1] = col_dir * spacing[0]
148    affine[:3, 2] = row_dir * spacing[1]
149    affine[:3, 3] = positions[np.argmin(projections)]
150    return labels, affine
151
152
153def _resample_labels(labels, affine, target_shape, target_affine):
154    """Resample a label volume onto the voxel grid of a reference image with nearest neighbor interpolation.
155
156    This is exact if both volumes are stored on the same grid (e.g. a DICOM-SEG object stored on a cropped
157    grid of the reference image) and downsamples segmentations that are stored on a finer grid.
158    """
159    to_label_index = np.linalg.inv(affine) @ target_affine
160    resampled = np.zeros(target_shape, dtype=labels.dtype)
161    yy, xx = np.meshgrid(np.arange(target_shape[1]), np.arange(target_shape[2]), indexing="ij")
162    for z in range(target_shape[0]):
163        target_indices = np.stack([np.full(yy.size, z), yy.ravel(), xx.ravel(), np.ones(yy.size)])
164        indices = np.round(to_label_index[:3] @ target_indices).astype("int")
165        valid = np.all((indices >= 0) & (indices < np.array(labels.shape)[:, None]), axis=0)
166        resampled[z].flat[valid] = labels[tuple(indices[:, valid])]
167    return resampled
168
169
170def _download_series(series_uids, dicom_dir):
171    """Download DICOM series from TCIA via the NBIA REST API into '<dicom_dir>/<SeriesInstanceUID>/'."""
172    os.makedirs(dicom_dir, exist_ok=True)
173    for uid in tqdm(series_uids, desc=f"Download {len(series_uids)} series from TCIA to {dicom_dir}"):
174        series_dir = os.path.join(dicom_dir, uid)
175        if os.path.exists(series_dir):  # This series has been downloaded already.
176            continue
177
178        # The series is downloaded as a zip archive, which is extracted to a temporary folder
179        # and only moved to the final location once it is complete.
180        with tempfile.TemporaryDirectory(dir=dicom_dir) as tmp_dir:
181            zip_path = os.path.join(tmp_dir, "series.zip")
182            with requests.get(util.NBIA_API_URL + "getImage", params={"SeriesInstanceUID": uid}, stream=True) as r:
183                r.raise_for_status()
184                with open(zip_path, "wb") as f:
185                    copyfileobj(r.raw, f)
186            tmp_series_dir = os.path.join(tmp_dir, "series")
187            util.unzip(zip_path, tmp_series_dir)
188            os.rename(tmp_series_dir, series_dir)
189
190
191def _get_series_metadata(csv_path, download):
192    """Fetch the metadata of all series in the ReMIND collection from TCIA and pair each segmentation
193    with the MRI series it references (via the study and the series description, e.g.
194    'tumor seg - MR ref: 3D_AX_T1_postcontrast').
195    """
196    if not os.path.exists(csv_path):
197        if not download:
198            raise RuntimeError(f"Cannot find the data at {csv_path}, but download was set to False.")
199        response = requests.get(util.NBIA_API_URL + "getSeries", params={"Collection": "ReMIND"})
200        response.raise_for_status()
201        metadata = response.json()
202        with open(csv_path, "w", newline="") as f:
203            writer = csv.DictWriter(f, fieldnames=sorted({key for row in metadata for key in row.keys()}))
204            writer.writeheader()
205            writer.writerows(metadata)
206
207    with open(csv_path, "r") as f:
208        metadata = list(csv.DictReader(f))
209
210    mr_series = {}
211    for row in metadata:
212        if row["Modality"] == "MR":
213            mr_series[(row["StudyInstanceUID"], row["SeriesDescription"])] = row
214
215    segmentations = {}
216    for row in metadata:
217        if row["Modality"] != "SEG":
218            continue
219        reference = row["SeriesDescription"].split("MR ref:")[-1].strip()
220        mr_row = mr_series[(row["StudyInstanceUID"], reference)]
221        segmentations.setdefault(mr_row["SeriesInstanceUID"], (mr_row, []))[1].append(row)
222
223    return segmentations
224
225
226def _preprocess_remind(dicom_dir, segmentations, preprocessed_dir):
227    import h5py
228
229    os.makedirs(preprocessed_dir, exist_ok=True)
230    for mr_uid, (mr_row, seg_rows) in tqdm(sorted(segmentations.items()), desc="Preprocess ReMIND"):
231        study = mr_row["StudyDesc"].lower()
232        modality = _get_modality(mr_row["SeriesDescription"])
233        fname = f"{mr_row['PatientID']}_{study}_{modality}_{mr_row['SeriesNumber']}.h5"
234        out_path = os.path.join(preprocessed_dir, fname)
235        if os.path.exists(out_path):
236            continue
237
238        volume, affine = _load_dicom_volume(os.path.join(dicom_dir, mr_uid))
239
240        # Each structure is stored in a separate DICOM-SEG object (possibly on a cropped or finer grid), which is
241        # resampled onto the MRI grid. The structures are written in the order of their label ids, so that more
242        # specific structures (e.g. the tumor) overwrite larger structures (e.g. the cerebrum) where they overlap.
243        labels = np.zeros(volume.shape, dtype="uint8")
244        structures = {row["SeriesDescription"].split(" seg")[0]: row["SeriesInstanceUID"] for row in seg_rows}
245        for structure, seg_uid in sorted(structures.items(), key=lambda item: LABEL_IDS[item[0]]):
246            seg_labels, seg_affine = _load_dicom_seg(glob(os.path.join(dicom_dir, seg_uid, "*.dcm"))[0])
247            mask = _resample_labels(seg_labels, seg_affine, volume.shape, affine) > 0
248            labels[mask] = LABEL_IDS[structure]
249
250        with h5py.File(out_path, "w") as f:
251            f.create_dataset("raw", data=volume, compression="gzip")
252            f.create_dataset("labels", data=labels, compression="gzip")
253            f.attrs["series_description"] = mr_row["SeriesDescription"]
254            f.attrs["structures"] = sorted(structures.keys())
255
256
257def get_remind_data(path: Union[os.PathLike, str], download: bool = False) -> str:
258    """Download the ReMIND dataset.
259
260    Args:
261        path: Filepath to a folder where the data is downloaded for further processing.
262        download: Whether to download the data if it is not present.
263
264    Returns:
265        Filepath where the preprocessed data is stored.
266    """
267    preprocessed_dir = os.path.join(path, "preprocessed")
268    if os.path.exists(preprocessed_dir):
269        return preprocessed_dir
270
271    os.makedirs(path, exist_ok=True)
272
273    # Download the manifest for reference and the series metadata of the collection. Only the MRI series
274    # with segmentations and the corresponding DICOM-SEG series are downloaded.
275    manifest_path = os.path.join(path, "ReMIND-Manifest-Sept-2023.tcia")
276    util.download_source(path=manifest_path, url=URL, download=download, checksum=CHECKSUM)
277    segmentations = _get_series_metadata(os.path.join(path, "remind_series.csv"), download)
278
279    dicom_dir = os.path.join(path, "dicom")
280    series_uids = sorted(segmentations.keys())
281    series_uids += sorted(row["SeriesInstanceUID"] for _, seg_rows in segmentations.values() for row in seg_rows)
282    if not all(os.path.exists(os.path.join(dicom_dir, uid)) for uid in series_uids):
283        if not download:
284            raise RuntimeError(f"Cannot find the data at {dicom_dir}, but download was set to False.")
285        _download_series(series_uids, dicom_dir)
286
287    _preprocess_remind(dicom_dir, segmentations, preprocessed_dir)
288    return preprocessed_dir
289
290
291def get_remind_paths(
292    path: Union[os.PathLike, str],
293    modality: Optional[Literal["t1c", "t1", "t2", "flair"]] = None,
294    study: Optional[Literal["preop", "intraop"]] = None,
295    download: bool = False,
296) -> List[str]:
297    """Get paths to the ReMIND data.
298
299    Args:
300        path: Filepath to a folder where the data is downloaded for further processing.
301        modality: The MRI sequence. One of 't1c', 't1', 't2' or 'flair'. If None, all sequences are returned.
302        study: The choice of study. Either 'preop' or 'intraop'. If None, both studies are returned.
303        download: Whether to download the data if it is not present.
304
305    Returns:
306        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
307    """
308    data_dir = get_remind_data(path, download)
309
310    if modality is not None and modality not in MODALITIES:
311        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {MODALITIES}.")
312    if study is not None and study not in ["preop", "intraop"]:
313        raise ValueError(f"'{study}' is not a valid study. Please choose one of 'preop' or 'intraop'.")
314
315    pattern = f"*_{'*' if study is None else study}_{'*' if modality is None else modality}_*.h5"
316    volume_paths = natsorted(glob(os.path.join(data_dir, pattern)))
317    return volume_paths
318
319
320def get_remind_dataset(
321    path: Union[os.PathLike, str],
322    patch_shape: Tuple[int, ...],
323    modality: Optional[Literal["t1c", "t1", "t2", "flair"]] = None,
324    study: Optional[Literal["preop", "intraop"]] = None,
325    resize_inputs: bool = False,
326    download: bool = False,
327    **kwargs
328) -> Dataset:
329    """Get the ReMIND dataset for brain tumor and resection segmentation.
330
331    Args:
332        path: Filepath to a folder where the data is downloaded for further processing.
333        patch_shape: The patch shape to use for training.
334        modality: The MRI sequence. One of 't1c', 't1', 't2' or 'flair'. If None, all sequences are returned.
335        study: The choice of study. Either 'preop' or 'intraop'. If None, both studies are returned.
336        resize_inputs: Whether to resize inputs to the desired patch shape.
337        download: Whether to download the data if it is not present.
338        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
339
340    Returns:
341        The segmentation dataset.
342    """
343    volume_paths = get_remind_paths(path, modality, study, download)
344
345    if resize_inputs:
346        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
347        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
348            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
349        )
350
351    return torch_em.default_segmentation_dataset(
352        raw_paths=volume_paths,
353        raw_key="raw",
354        label_paths=volume_paths,
355        label_key="labels",
356        patch_shape=patch_shape,
357        is_seg_dataset=True,
358        **kwargs
359    )
360
361
362def get_remind_loader(
363    path: Union[os.PathLike, str],
364    batch_size: int,
365    patch_shape: Tuple[int, ...],
366    modality: Optional[Literal["t1c", "t1", "t2", "flair"]] = None,
367    study: Optional[Literal["preop", "intraop"]] = None,
368    resize_inputs: bool = False,
369    download: bool = False,
370    **kwargs
371) -> DataLoader:
372    """Get the ReMIND dataloader for brain tumor and resection segmentation.
373
374    Args:
375        path: Filepath to a folder where the data is downloaded for further processing.
376        batch_size: The batch size for training.
377        patch_shape: The patch shape to use for training.
378        modality: The MRI sequence. One of 't1c', 't1', 't2' or 'flair'. If None, all sequences are returned.
379        study: The choice of study. Either 'preop' or 'intraop'. If None, both studies are returned.
380        resize_inputs: Whether to resize inputs to the desired patch shape.
381        download: Whether to download the data if it is not present.
382        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
383
384    Returns:
385        The DataLoader.
386    """
387    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
388    dataset = get_remind_dataset(path, patch_shape, modality, study, resize_inputs, download, **ds_kwargs)
389    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)