torch_em.data.datasets.medical.ultracortex

  1"""The UltraCortex dataset contains submillimeter ultra-high field 9.4T brain MR images with manual
  2cortical gray and white matter segmentations.
  3
  4The dataset consists of 86 structural MR images (0.6-0.8mm resolution, MP-RAGE and MP2RAGE sequences),
  5of which 12 have manual cortical segmentations into gray and white matter, independently validated by
  6two expert neuroradiologists. This module only exposes the 12 volumes with manual segmentations.
  7
  8The dataset is located at https://openneuro.org/datasets/ds005216, released under the CC0 1.0 license.
  9
 10The dataset is from the publication https://doi.org/10.1038/s41597-025-04779-2.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 17from typing import Union, Tuple, List
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26BASE_URL = "https://s3.amazonaws.com/openneuro.org/ds005216"
 27
 28# The 12 subjects with manual cortical gray / white matter segmentations, out of the 86 total subjects.
 29SUBJECT_IDS = [3, 9, 20, 29, 37, 44, 45, 46, 57, 69, 70, 73]
 30
 31
 32def get_ultracortex_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the UltraCortex dataset.
 34
 35    Args:
 36        path: Filepath to a folder where the data is downloaded for further processing.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath where the data is downloaded.
 41    """
 42    os.makedirs(path, exist_ok=True)
 43
 44    for sub_id in SUBJECT_IDS:
 45        raw_path = os.path.join(path, f"sub-{sub_id}_ses-1_T1w.nii")
 46        raw_url = f"{BASE_URL}/sub-{sub_id}/ses-1/anat/sub-{sub_id}_ses-1_T1w.nii"
 47        util.download_source(path=raw_path, url=raw_url, download=download)
 48
 49        label_path = os.path.join(path, f"sub-{sub_id}_ses-1_seg.nii")
 50        label_url = f"{BASE_URL}/derivatives/manual_segmentation/sub-{sub_id}_ses-1_seg.nii"
 51        util.download_source(path=label_path, url=label_url, download=download)
 52
 53    return path
 54
 55
 56def get_ultracortex_paths(
 57    path: Union[os.PathLike, str], download: bool = False
 58) -> Tuple[List[str], List[str]]:
 59    """Get paths to the UltraCortex data.
 60
 61    Args:
 62        path: Filepath to a folder where the data is downloaded for further processing.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    get_ultracortex_data(path, download)
 70
 71    label_paths = natsorted(glob(os.path.join(path, "sub-*_ses-1_seg.nii")))
 72    raw_paths = [p.replace("_seg.nii", "_T1w.nii") for p in label_paths]
 73    assert all(os.path.exists(p) for p in raw_paths), "Some image volumes are missing."
 74    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 75
 76    return raw_paths, label_paths
 77
 78
 79def get_ultracortex_dataset(
 80    path: Union[os.PathLike, str],
 81    patch_shape: Tuple[int, ...],
 82    resize_inputs: bool = False,
 83    download: bool = False,
 84    **kwargs
 85) -> Dataset:
 86    """Get the UltraCortex dataset for cortical gray and white matter segmentation in 9.4T brain MRI.
 87
 88    Args:
 89        path: Filepath to a folder where the data is downloaded for further processing.
 90        patch_shape: The patch shape to use for training.
 91        resize_inputs: Whether to resize the inputs to the patch shape.
 92        download: Whether to download the data if it is not present.
 93        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 94
 95    Returns:
 96        The segmentation dataset.
 97    """
 98    raw_paths, label_paths = get_ultracortex_paths(path, download)
 99
100    if resize_inputs:
101        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
102        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
103            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
104        )
105
106    return torch_em.default_segmentation_dataset(
107        raw_paths=raw_paths,
108        raw_key="data",
109        label_paths=label_paths,
110        label_key="data",
111        patch_shape=patch_shape,
112        is_seg_dataset=True,
113        **kwargs
114    )
115
116
117def get_ultracortex_loader(
118    path: Union[os.PathLike, str],
119    batch_size: int,
120    patch_shape: Tuple[int, ...],
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> DataLoader:
125    """Get the UltraCortex dataloader for cortical gray and white matter segmentation in 9.4T brain MRI.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        batch_size: The batch size for training.
130        patch_shape: The patch shape to use for training.
131        resize_inputs: Whether to resize the inputs to the patch shape.
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
134
135    Returns:
136        The DataLoader.
137    """
138    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
139    dataset = get_ultracortex_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
140    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)