torch_em.data.datasets.medical.mendeley_ms

  1"""The Mendeley MS dataset contains annotations for multiple sclerosis lesion segmentation in brain MRI.
  2
  3It comprises T1-weighted, T2-weighted and T2-FLAIR scans of 60 MS patients with consensus manual lesion
  4segmentations for each of the three modalities (the scans of a patient are not co-registered and have
  5different shapes, so the labels are provided per modality). This is the 'MS Lesion' dataset of the
  6RadioActive benchmark (https://arxiv.org/abs/2411.07885), which uses the FLAIR scans.
  7The label ids are: 0 = background, 1 = MS lesion.
  8
  9The dataset is located at https://data.mendeley.com/datasets/8bctsm8jz7/1 (CC BY 4.0).
 10
 11This dataset is from the publication https://doi.org/10.1016/j.dib.2022.108139.
 12Please cite it if you use this dataset in your research.
 13"""
 14
 15import os
 16from glob import glob
 17from natsort import natsorted
 18from typing import Union, Tuple, Literal, List
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://data.mendeley.com/public-files/datasets/8bctsm8jz7/files/9356efeb-dcd8-4213-a2d4-8febe9f1a5db/file_downloaded"  # noqa
 28CHECKSUM = "c90f0f47c9e1a5e0fafc87b77dfcbcb09ac0cf9ffdaa333aaec1b9c63d31a7b3"
 29
 30LABEL_IDS = {"background": 0, "ms_lesion": 1}
 31
 32MODALITIES = {"flair": "Flair", "t1": "T1", "t2": "T2"}
 33
 34
 35def get_mendeley_ms_data(path: Union[os.PathLike, str], download: bool = False):
 36    """Download the Mendeley MS dataset.
 37
 38    Args:
 39        path: Filepath to a folder where the data is downloaded for further processing.
 40        download: Whether to download the data if it is not present.
 41    """
 42    if len(glob(os.path.join(path, "Patient-*"))) == 60:
 43        return
 44
 45    os.makedirs(path, exist_ok=True)
 46
 47    zip_path = os.path.join(path, "ms_brain_mri.zip")
 48    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 49    util.unzip(zip_path=zip_path, dst=path)
 50
 51
 52def get_mendeley_ms_paths(
 53    path: Union[os.PathLike, str], modality: Literal["flair", "t1", "t2"] = "flair", download: bool = False
 54) -> Tuple[List[str], List[str]]:
 55    """Get paths to the Mendeley MS data.
 56
 57    Args:
 58        path: Filepath to a folder where the data is downloaded for further processing.
 59        modality: The MRI modality. Either 'flair', 't1' or 't2'.
 60        download: Whether to download the data if it is not present.
 61
 62    Returns:
 63        List of filepaths for the image data.
 64        List of filepaths for the label data.
 65    """
 66    get_mendeley_ms_data(path, download)
 67
 68    if modality not in MODALITIES:
 69        raise ValueError(f"'{modality}' is not a valid modality. Choose one of {list(MODALITIES)}.")
 70    modality = MODALITIES[modality]
 71
 72    raw_paths = natsorted(glob(os.path.join(path, "Patient-*", f"*-{modality}.nii")))
 73    raw_paths = [p for p in raw_paths if "LesionSeg" not in os.path.basename(p)]
 74    label_paths = [p.replace(f"-{modality}.nii", f"-LesionSeg-{modality}.nii") for p in raw_paths]
 75    assert all(os.path.exists(p) for p in label_paths)
 76
 77    return raw_paths, label_paths
 78
 79
 80def get_mendeley_ms_dataset(
 81    path: Union[os.PathLike, str],
 82    patch_shape: Tuple[int, ...],
 83    modality: Literal["flair", "t1", "t2"] = "flair",
 84    resize_inputs: bool = False,
 85    download: bool = False,
 86    **kwargs
 87) -> Dataset:
 88    """Get the Mendeley MS dataset for MS lesion segmentation.
 89
 90    Args:
 91        path: Filepath to a folder where the data is downloaded for further processing.
 92        patch_shape: The patch shape to use for training.
 93        modality: The MRI modality. Either 'flair', 't1' or 't2'.
 94        resize_inputs: Whether to resize inputs to the desired patch shape.
 95        download: Whether to download the data if it is not present.
 96        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 97
 98    Returns:
 99        The segmentation dataset.
100    """
101    raw_paths, label_paths = get_mendeley_ms_paths(path, modality, download)
102
103    if resize_inputs:
104        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
105        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
106            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
107        )
108
109    return torch_em.default_segmentation_dataset(
110        raw_paths=raw_paths,
111        raw_key="data",
112        label_paths=label_paths,
113        label_key="data",
114        patch_shape=patch_shape,
115        is_seg_dataset=True,
116        **kwargs
117    )
118
119
120def get_mendeley_ms_loader(
121    path: Union[os.PathLike, str],
122    batch_size: int,
123    patch_shape: Tuple[int, ...],
124    modality: Literal["flair", "t1", "t2"] = "flair",
125    resize_inputs: bool = False,
126    download: bool = False,
127    **kwargs
128) -> DataLoader:
129    """Get the Mendeley MS dataloader for MS lesion segmentation.
130
131    Args:
132        path: Filepath to a folder where the data is downloaded for further processing.
133        batch_size: The batch size for training.
134        patch_shape: The patch shape to use for training.
135        modality: The MRI modality. Either 'flair', 't1' or 't2'.
136        resize_inputs: Whether to resize inputs to the desired patch shape.
137        download: Whether to download the data if it is not present.
138        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
139
140    Returns:
141        The DataLoader.
142    """
143    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
144    dataset = get_mendeley_ms_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
145    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)