torch_em.data.datasets.medical.mcrib

  1"""The M-CRIB dataset contains cortical and subcortical parcellations of neonatal brain MRI.
  2
  3The dataset consists of 10 healthy term-born neonates (scanned at 40-43 weeks gestational age) with T2-weighted
  4and T1-weighted MRI (the T1 volumes are provided registered to the T2 volumes) and manual parcellations
  5following the M-CRIB 2.0 protocol, which is compatible with the adult Desikan-Killiany cortical atlas
  6and the FreeSurfer subcortical labels.
  7
  8NOTE: The label volumes are semantic parcellations with FreeSurfer-style ids (94 structures per volume):
  9- subcortical structures use the FreeSurfer ids (e.g. 2: Left-Cerebral-White-Matter, 4: Left-Lateral-Ventricle,
 10  9: Left-Thalamus, 17: Left-Hippocampus, 24: CSF, 41: Right-Cerebral-White-Matter, 170: brainstem,
 11  192: Corpus_Callosum, and the cerebellar labels 75, 76, 90, 91, 93),
 12- the left hemisphere cortical regions use the ids 1000-1035 (e.g. 1002: ctx-lh-caudalanteriorcingulate),
 13- the right hemisphere cortical regions use the ids 2000-2035 (e.g. 2002: ctx-rh-caudalanteriorcingulate).
 14The complete lookup table (id, RGB color, name) is shipped with the data in
 15'M-CRIB_2-0_labels_itk_format.txt', which is downloaded next to the volumes.
 16
 17The dataset is located at https://osf.io/4vthr/.
 18
 19This dataset is from the publications https://doi.org/10.1038/sdata.2017.57 (M-CRIB) and
 20https://doi.org/10.3389/fnins.2019.00034 (M-CRIB 2.0).
 21Please cite them if you use this dataset in your research.
 22"""
 23
 24import os
 25from glob import glob
 26from natsort import natsorted
 27from typing import Union, Tuple, Literal, List
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36# The folder zips are generated on-the-fly by OSF, hence the checksums of the archives are not reliable.
 37URLS = {
 38    "T2": "https://files.osf.io/v1/resources/4vthr/providers/osfstorage/5d36b7efa667db0019f9dbc9/?zip=",
 39    "T1": "https://files.osf.io/v1/resources/4vthr/providers/osfstorage/5d37bbe0a667db0018fc7ab0/?zip=",
 40    "labels": "https://files.osf.io/v1/resources/4vthr/providers/osfstorage/5d36b267251f0e0017091695/?zip=",
 41    "lookup_table": "https://osf.io/download/9u42m/",
 42}
 43
 44CHECKSUMS = {
 45    "T2": None,
 46    "T1": None,
 47    "labels": None,
 48    "lookup_table": "bcd373c68399220a8884606b4abfd0df474edcb7522c1b292fa1b60f378b0505",
 49}
 50
 51
 52def get_mcrib_data(path: Union[os.PathLike, str], modality: Literal["T2", "T1"] = "T2", download: bool = False):
 53    """Download the M-CRIB dataset.
 54
 55    Args:
 56        path: Filepath to a folder where the data is downloaded for further processing.
 57        modality: The MRI modality. Either 'T2' or 'T1' (the T1 volumes registered to the T2 volumes).
 58        download: Whether to download the data if it is not present.
 59    """
 60    if modality not in ["T2", "T1"]:
 61        raise ValueError(f"'{modality}' is not a valid modality. Choose either 'T2' or 'T1'.")
 62
 63    os.makedirs(path, exist_ok=True)
 64
 65    for name in [modality, "labels"]:
 66        data_dir = os.path.join(path, name)
 67        if os.path.exists(data_dir):
 68            continue
 69
 70        zip_path = os.path.join(path, f"{name}.zip")
 71        util.download_source(path=zip_path, url=URLS[name], download=download, checksum=CHECKSUMS[name])
 72        util.unzip(zip_path=zip_path, dst=data_dir)
 73
 74    lut_path = os.path.join(path, "M-CRIB_2-0_labels_itk_format.txt")
 75    util.download_source(path=lut_path, url=URLS["lookup_table"], download=download, checksum=CHECKSUMS["lookup_table"])
 76
 77
 78def get_mcrib_paths(
 79    path: Union[os.PathLike, str], modality: Literal["T2", "T1"] = "T2", download: bool = False
 80) -> Tuple[List[str], List[str]]:
 81    """Get paths to the M-CRIB data.
 82
 83    Args:
 84        path: Filepath to a folder where the data is downloaded for further processing.
 85        modality: The MRI modality. Either 'T2' or 'T1' (the T1 volumes registered to the T2 volumes).
 86        download: Whether to download the data if it is not present.
 87
 88    Returns:
 89        List of filepaths for the image data.
 90        List of filepaths for the label data.
 91    """
 92    get_mcrib_data(path, modality, download)
 93
 94    label_paths = natsorted(glob(os.path.join(path, "labels", "M-CRIB_2-0_P*_parc.nii.gz")))
 95    suffix = "T2" if modality == "T2" else "T1_registered_to_T2"
 96    raw_paths = [
 97        os.path.join(path, modality, os.path.basename(p).replace("_2-0", "").replace("parc", suffix))
 98        for p in label_paths
 99    ]
100    assert all(os.path.exists(p) for p in raw_paths), "Some image volumes are missing."
101    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
102
103    return raw_paths, label_paths
104
105
106def get_mcrib_dataset(
107    path: Union[os.PathLike, str],
108    patch_shape: Tuple[int, ...],
109    modality: Literal["T2", "T1"] = "T2",
110    resize_inputs: bool = False,
111    download: bool = False,
112    **kwargs
113) -> Dataset:
114    """Get the M-CRIB dataset for neonatal brain parcellation.
115
116    Args:
117        path: Filepath to a folder where the data is downloaded for further processing.
118        patch_shape: The patch shape to use for training.
119        modality: The MRI modality. Either 'T2' or 'T1' (the T1 volumes registered to the T2 volumes).
120        resize_inputs: Whether to resize inputs to the desired patch shape.
121        download: Whether to download the data if it is not present.
122        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
123
124    Returns:
125        The segmentation dataset.
126    """
127    raw_paths, label_paths = get_mcrib_paths(path, modality, download)
128
129    if resize_inputs:
130        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
131        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
132            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
133        )
134
135    return torch_em.default_segmentation_dataset(
136        raw_paths=raw_paths,
137        raw_key="data",
138        label_paths=label_paths,
139        label_key="data",
140        patch_shape=patch_shape,
141        is_seg_dataset=True,
142        **kwargs
143    )
144
145
146def get_mcrib_loader(
147    path: Union[os.PathLike, str],
148    batch_size: int,
149    patch_shape: Tuple[int, ...],
150    modality: Literal["T2", "T1"] = "T2",
151    resize_inputs: bool = False,
152    download: bool = False,
153    **kwargs
154) -> DataLoader:
155    """Get the M-CRIB dataloader for neonatal brain parcellation.
156
157    Args:
158        path: Filepath to a folder where the data is downloaded for further processing.
159        batch_size: The batch size for training.
160        patch_shape: The patch shape to use for training.
161        modality: The MRI modality. Either 'T2' or 'T1' (the T1 volumes registered to the T2 volumes).
162        resize_inputs: Whether to resize inputs to the desired patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
165
166    Returns:
167        The DataLoader.
168    """
169    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
170    dataset = get_mcrib_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
171    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)