torch_em.data.datasets.medical.mrbrains18

  1"""The MRBrainS18 dataset contains annotations for the segmentation of brain structures and of
  2white matter lesions in multi-sequence brain MRI.
  3
  4The data was curated for the MRBrainS18 challenge (https://mrbrains18.isi.uu.nl), which was held at MICCAI
  52018 and is the successor of the MRBrainS13 challenge. It consists of 30 subjects scanned at the UMC Utrecht
  6with a 3T scanner, 7 of them released as the training set and the remaining 23 as the test set. The reference
  7standard was released for the test subjects as well, so this module provides 30 annotated volumes over the
  8two official splits.
  9
 10Three co-registered sequences are available per subject and can be selected with the 'modality' argument:
 11a 3D T1-weighted scan ('t1', registered to the label grid), a multi-slice T1-weighted inversion recovery scan
 12('ir') and a multi-slice T2 FLAIR scan ('flair'). All of them are bias field corrected with N4ITK and have a
 13voxel size of 0.958 x 0.958 x 3.0 mm.
 14
 15The label ids are described in `LABEL_IDS`: 0 = background, 1 = cortical grey matter, 2 = basal ganglia,
 163 = white matter, 4 = white matter lesions, 5 = cerebrospinal fluid in the extracerebral space,
 176 = ventricles, 7 = cerebellum, 8 = brain stem, 9 = infarction, 10 = other.
 18NOTE: The official challenge evaluation only scores the ids 1 to 8; the ids 9 and 10 are excluded from it.
 19
 20The three sequences and the reference standard of a subject are bundled into one hdf5 file per subject by
 21this module, with the slice axis first (the keys are 'raw/t1', 'raw/ir', 'raw/flair' and 'labels').
 22
 23The data is hosted at https://doi.org/10.34894/E0U32Q and is free to download, but it may only be used
 24under the terms of the UMC Utrecht license that is distributed with it.
 25
 26This dataset is from the publication https://doi.org/10.3389/fncom.2019.00093.
 27Please cite it if you use this dataset in your research.
 28"""
 29
 30import os
 31from glob import glob
 32from tqdm import tqdm
 33from natsort import natsorted
 34from typing import Union, Tuple, List, Optional, Literal
 35
 36import numpy as np
 37
 38from torch.utils.data import Dataset, DataLoader
 39
 40import torch_em
 41
 42from .. import util
 43
 44
 45URLS = {
 46    "train": "https://dataverse.nl/api/access/datafile/402709",
 47    "test": "https://dataverse.nl/api/access/datafile/402708",
 48}
 49
 50CHECKSUMS = {
 51    "train": "c29b133c0d0e7486563d3a92f09e3ae67986880130bb8a91096fc1fdf0c93e95",
 52    "test": "1576500224bd98ef66fa0277af13bd304ed491974a0de96006f6779e5363d3dd",
 53}
 54
 55LABEL_IDS = {
 56    "background": 0,
 57    "cortical_grey_matter": 1,
 58    "basal_ganglia": 2,
 59    "white_matter": 3,
 60    "white_matter_lesions": 4,
 61    "csf_extracerebral": 5,
 62    "ventricles": 6,
 63    "cerebellum": 7,
 64    "brain_stem": 8,
 65    "infarction": 9,
 66    "other": 10,
 67}
 68
 69SPLITS = {"train": "training", "test": "test"}
 70
 71MODALITIES = {"t1": "reg_T1.nii.gz", "ir": "reg_IR.nii.gz", "flair": "FLAIR.nii.gz"}
 72
 73SUBJECT_IDS = {
 74    "train": [1, 4, 5, 7, 14, 27, 29],
 75    "test": [2, 3, 6, 8, 9, 10, 11, 12, 13, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 28, 30],
 76}
 77
 78
 79def _preprocess_inputs(data_dir, subject_ids, preprocessed_dir):
 80    import h5py
 81    import nibabel as nib
 82
 83    os.makedirs(preprocessed_dir, exist_ok=True)
 84    for subject_id in tqdm(subject_ids, desc="Preprocessing the MRBrainS18 subjects"):
 85        volume_path = os.path.join(preprocessed_dir, f"subject_{subject_id:02}.h5")
 86        if os.path.exists(volume_path):
 87            continue
 88
 89        subject_dir = os.path.join(data_dir, str(subject_id))
 90
 91        # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
 92        labels = np.asarray(nib.load(os.path.join(subject_dir, "segm.nii.gz")).dataobj).T
 93
 94        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
 95        with h5py.File(f"{volume_path}.tmp", "w") as f:
 96            for modality, fname in MODALITIES.items():
 97                raw = np.asarray(nib.load(os.path.join(subject_dir, "pre", fname)).dataobj).T
 98                f.create_dataset(f"raw/{modality}", data=raw, compression="gzip")
 99
100            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
101
102        os.rename(f"{volume_path}.tmp", volume_path)
103
104
105def get_mrbrains18_data(
106    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
107) -> str:
108    """Download the MRBrainS18 dataset.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        split: The choice of data split. Either 'train' or 'test'.
113        download: Whether to download the data if it is not present.
114
115    Returns:
116        Filepath where the preprocessed data is stored.
117    """
118    if split not in SPLITS:
119        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(SPLITS.keys())}.")
120
121    subject_ids = SUBJECT_IDS[split]
122    preprocessed_dir = os.path.join(path, "preprocessed", split)
123    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == len(subject_ids):
124        return preprocessed_dir
125
126    os.makedirs(path, exist_ok=True)
127
128    data_dir = os.path.join(path, SPLITS[split])
129    if not os.path.exists(data_dir):
130        zip_path = os.path.join(path, f"{SPLITS[split]}.zip")
131        util.download_source(path=zip_path, url=URLS[split], download=download, checksum=CHECKSUMS[split])
132        util.unzip(zip_path=zip_path, dst=path)
133
134    _preprocess_inputs(data_dir, subject_ids, preprocessed_dir)
135    return preprocessed_dir
136
137
138def get_mrbrains18_paths(
139    path: Union[os.PathLike, str],
140    split: Optional[Literal["train", "test"]] = None,
141    download: bool = False,
142) -> List[str]:
143    """Get paths to the MRBrainS18 data.
144
145    Args:
146        path: Filepath to a folder where the data is downloaded for further processing.
147        split: The choice of data split. Either 'train' or 'test'. If None, all subjects are returned.
148        download: Whether to download the data if it is not present.
149
150    Returns:
151        List of filepaths for the hdf5 files, which contain the image data ('raw/<modality>')
152        and the label data ('labels').
153    """
154    splits = list(SPLITS.keys()) if split is None else [split]
155    volume_paths = []
156    for curr_split in splits:
157        volume_paths.extend(natsorted(glob(os.path.join(get_mrbrains18_data(path, curr_split, download), "*.h5"))))
158
159    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{path}'."
160    return volume_paths
161
162
163def get_mrbrains18_dataset(
164    path: Union[os.PathLike, str],
165    patch_shape: Tuple[int, ...],
166    split: Optional[Literal["train", "test"]] = None,
167    modality: Literal["t1", "ir", "flair"] = "t1",
168    resize_inputs: bool = False,
169    download: bool = False,
170    **kwargs
171) -> Dataset:
172    """Get the MRBrainS18 dataset for brain structure and white matter lesion segmentation.
173
174    Args:
175        path: Filepath to a folder where the data is downloaded for further processing.
176        patch_shape: The patch shape to use for training.
177        split: The choice of data split. Either 'train' or 'test'. If None, all subjects are returned.
178        modality: The MRI sequence. Either 't1', 'ir' or 'flair'.
179        resize_inputs: Whether to resize inputs to the desired patch shape.
180        download: Whether to download the data if it is not present.
181        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
182
183    Returns:
184        The segmentation dataset.
185    """
186    if modality not in MODALITIES:
187        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {list(MODALITIES.keys())}.")
188
189    volume_paths = get_mrbrains18_paths(path, split, download)
190
191    if resize_inputs:
192        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
193        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
194            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
195        )
196
197    return torch_em.default_segmentation_dataset(
198        raw_paths=volume_paths,
199        raw_key=f"raw/{modality}",
200        label_paths=volume_paths,
201        label_key="labels",
202        patch_shape=patch_shape,
203        is_seg_dataset=True,
204        **kwargs
205    )
206
207
208def get_mrbrains18_loader(
209    path: Union[os.PathLike, str],
210    batch_size: int,
211    patch_shape: Tuple[int, ...],
212    split: Optional[Literal["train", "test"]] = None,
213    modality: Literal["t1", "ir", "flair"] = "t1",
214    resize_inputs: bool = False,
215    download: bool = False,
216    **kwargs
217) -> DataLoader:
218    """Get the MRBrainS18 dataloader for brain structure and white matter lesion 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        split: The choice of data split. Either 'train' or 'test'. If None, all subjects are returned.
225        modality: The MRI sequence. Either 't1', 'ir' or 'flair'.
226        resize_inputs: Whether to resize inputs to the desired patch shape.
227        download: Whether to download the data if it is not present.
228        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
229
230    Returns:
231        The DataLoader.
232    """
233    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
234    dataset = get_mrbrains18_dataset(path, patch_shape, split, modality, resize_inputs, download, **ds_kwargs)
235    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)