torch_em.data.datasets.medical.pedims

  1"""The PediMS dataset contains annotations for multiple sclerosis lesion segmentation in pediatric brain MRI.
  2
  3The dataset comprises 28 longitudinal MRI exams from 9 pediatric MS patients (1 to 6 timepoints each),
  4acquired with T1-weighted, T2-weighted and T2-FLAIR sequences. Each timepoint ships a consensus lesion
  5mask, delineated and validated by senior clinical experts, in the native FLAIR space (the T1 and T2
  6scans are provided in their own native spaces and are not registered to the lesion mask, so this module
  7only exposes the FLAIR scan, which is used by this dataset for lesion delineation).
  8The label ids are: 0 = background, 1 = MS lesion.
  9
 10The dataset is located at https://doi.org/10.6084/m9.figshare.28701065.v1 (CC BY 4.0).
 11
 12This dataset is from the publication https://doi.org/10.1038/s41597-025-05346-5.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://ndownloader.figshare.com/articles/28701065/versions/1"
 29CHECKSUM = "2bd6dd209654a79247ba6340cf39afca7e6d20beacb61fea186432d22b05122a"
 30
 31LABEL_IDS = {"background": 0, "ms_lesion": 1}
 32
 33
 34def get_pedims_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the PediMS dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the data is downloaded for further processing.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        Filepath where the data is downloaded.
 43    """
 44    data_dir = os.path.join(path, "PediMS")
 45    if os.path.exists(data_dir):
 46        return data_dir
 47
 48    os.makedirs(path, exist_ok=True)
 49
 50    zip_path = os.path.join(path, "pedims.zip")
 51    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 52    util.unzip(zip_path=zip_path, dst=path)
 53
 54    return data_dir
 55
 56
 57def get_pedims_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 58    """Get paths to the PediMS data.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        download: Whether to download the data if it is not present.
 63
 64    Returns:
 65        List of filepaths for the image data.
 66        List of filepaths for the label data.
 67    """
 68    data_dir = get_pedims_data(path, download)
 69
 70    raw_paths = natsorted(glob(os.path.join(data_dir, "P*", "T*", "processed", "brain_FLAIR.nii.gz")))
 71    label_paths = [p.replace("brain_FLAIR.nii.gz", "Consensus.nii") for p in raw_paths]
 72
 73    assert len(raw_paths) == 28, f"Expected 28 timepoints, found {len(raw_paths)} in '{data_dir}'."
 74    for label_path in label_paths:
 75        assert os.path.exists(label_path), label_path
 76
 77    return raw_paths, label_paths
 78
 79
 80def get_pedims_dataset(
 81    path: Union[os.PathLike, str],
 82    patch_shape: Tuple[int, ...],
 83    resize_inputs: bool = False,
 84    download: bool = False,
 85    **kwargs
 86) -> Dataset:
 87    """Get the PediMS dataset for pediatric multiple sclerosis lesion segmentation.
 88
 89    Args:
 90        path: Filepath to a folder where the data is downloaded for further processing.
 91        patch_shape: The patch shape to use for training.
 92        resize_inputs: Whether to resize inputs to the desired patch shape.
 93        download: Whether to download the data if it is not present.
 94        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 95
 96    Returns:
 97        The segmentation dataset.
 98    """
 99    raw_paths, label_paths = get_pedims_paths(path, download)
100
101    if resize_inputs:
102        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
103        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
104            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
105        )
106
107    return torch_em.default_segmentation_dataset(
108        raw_paths=raw_paths,
109        raw_key="data",
110        label_paths=label_paths,
111        label_key="data",
112        patch_shape=patch_shape,
113        is_seg_dataset=True,
114        **kwargs
115    )
116
117
118def get_pedims_loader(
119    path: Union[os.PathLike, str],
120    batch_size: int,
121    patch_shape: Tuple[int, ...],
122    resize_inputs: bool = False,
123    download: bool = False,
124    **kwargs
125) -> DataLoader:
126    """Get the PediMS dataloader for pediatric multiple sclerosis lesion segmentation.
127
128    Args:
129        path: Filepath to a folder where the data is downloaded for further processing.
130        batch_size: The batch size for training.
131        patch_shape: The patch shape to use for training.
132        resize_inputs: Whether to resize inputs to the desired patch shape.
133        download: Whether to download the data if it is not present.
134        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
135
136    Returns:
137        The DataLoader.
138    """
139    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
140    dataset = get_pedims_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
141    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)