torch_em.data.datasets.medical.totalsegmentator_mri

  1"""The TotalSegmentator MRI dataset contains annotations for 50 anatomical structures in MRI scans.
  2
  3The dataset (v3.0.0) consists of 1296 MRI volumes (whole body, pediatric and clinical routine scans with a wide
  4range of sequences, scanners and institutions) with an official train / test split (see 'meta.csv'; there is no
  5validation split).
  6Each anatomical structure is provided as a separate binary mask. `get_totalsegmentator_mri_data` merges these masks
  7into a single semantic label volume per case, where the label id of each structure is its (1-based) position in
  8`CLASS_NAMES` (see `CLASS_IDS` for the name -> id mapping). This is the class order of the 'total_mr' task in the
  9TotalSegmentator repository (https://github.com/wasserth/TotalSegmentator). The masks of a few structures may
 10overlap, in this case the structure with the higher label id takes precedence.
 11
 12The dataset is located at https://doi.org/10.5281/zenodo.11367005.
 13
 14This dataset is from the publication https://doi.org/10.1148/radiol.241613.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from typing import Union, Tuple, Literal, List, Optional
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27from .totalsegmentator import merge_all_segmentations, read_split
 28
 29
 30URL = "https://zenodo.org/records/22688334/files/TotalsegmentatorMRI_dataset_v300.zip"
 31CHECKSUM = "791696df98dbbc5136c19d6b23335779a9bdcda7fcc3842e885869756ca17589"
 32
 33CLASS_NAMES = [
 34    "spleen", "kidney_right", "kidney_left", "gallbladder", "liver", "stomach", "pancreas", "adrenal_gland_right",
 35    "adrenal_gland_left", "lung_left", "lung_right", "esophagus", "small_bowel", "duodenum", "colon",
 36    "urinary_bladder", "prostate", "sacrum", "vertebrae", "intervertebral_discs", "spinal_cord", "heart", "aorta",
 37    "inferior_vena_cava", "portal_vein_and_splenic_vein", "iliac_artery_left", "iliac_artery_right",
 38    "iliac_vena_left", "iliac_vena_right", "humerus_left", "humerus_right", "scapula_left", "scapula_right",
 39    "clavicula_left", "clavicula_right", "femur_left", "femur_right", "hip_left", "hip_right", "gluteus_maximus_left",
 40    "gluteus_maximus_right", "gluteus_medius_left", "gluteus_medius_right", "gluteus_minimus_left",
 41    "gluteus_minimus_right", "autochthon_left", "autochthon_right", "iliopsoas_left", "iliopsoas_right", "brain",
 42]
 43"""The anatomical structures of the TotalSegmentator MRI dataset. The label id of a structure is its 1-based index."""
 44
 45CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 46"""Mapping from the name of an anatomical structure to its label id in the merged label volumes."""
 47
 48
 49def get_totalsegmentator_mri_data(
 50    path: Union[os.PathLike, str], download: bool = False, n_workers: Optional[int] = None
 51) -> str:
 52    """Download the TotalSegmentator MRI dataset and merge the per-class masks into semantic label volumes.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        download: Whether to download the data if it is not present.
 57        n_workers: The number of parallel workers for merging the per-class masks.
 58
 59    Returns:
 60        Filepath where the data is downloaded.
 61    """
 62    data_dir = os.path.join(path, "Totalsegmentator_dataset_v300")
 63    if not os.path.exists(os.path.join(data_dir, "meta.csv")):
 64        os.makedirs(path, exist_ok=True)
 65        zip_path = os.path.join(path, "TotalsegmentatorMRI_dataset_v300.zip")
 66        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 67        util.unzip(zip_path=zip_path, dst=path)
 68
 69    case_dirs = sorted(glob(os.path.join(data_dir, "s*")))
 70    merge_all_segmentations(case_dirs, CLASS_NAMES, n_workers)
 71
 72    return data_dir
 73
 74
 75def get_totalsegmentator_mri_paths(
 76    path: Union[os.PathLike, str], split: Literal['train', 'test'], download: bool = False
 77) -> Tuple[List[str], List[str]]:
 78    """Get paths to the TotalSegmentator MRI data.
 79
 80    Args:
 81        path: Filepath to a folder where the data is downloaded for further processing.
 82        split: The choice of data split. Either 'train' or 'test'.
 83        download: Whether to download the data if it is not present.
 84
 85    Returns:
 86        List of filepaths for the image data.
 87        List of filepaths for the label data.
 88    """
 89    data_dir = get_totalsegmentator_mri_data(path, download)
 90    case_ids = read_split(os.path.join(data_dir, "meta.csv"), split, valid_splits=("train", "test"))
 91
 92    raw_paths = [os.path.join(data_dir, case_id, "mri.nii.gz") for case_id in case_ids]
 93    label_paths = [os.path.join(data_dir, case_id, "labels.nii.gz") for case_id in case_ids]
 94    assert all(os.path.exists(p) for p in raw_paths + label_paths)
 95
 96    return raw_paths, label_paths
 97
 98
 99def get_totalsegmentator_mri_dataset(
100    path: Union[os.PathLike, str],
101    patch_shape: Tuple[int, ...],
102    split: Literal['train', 'test'],
103    resize_inputs: bool = False,
104    download: bool = False,
105    **kwargs
106) -> Dataset:
107    """Get the TotalSegmentator MRI dataset for segmentation of anatomical structures in MRI.
108
109    Args:
110        path: Filepath to a folder where the data is downloaded for further processing.
111        patch_shape: The patch shape to use for training.
112        split: The choice of data split. Either 'train' or 'test'.
113        resize_inputs: Whether to resize inputs to the desired patch shape.
114        download: Whether to download the data if it is not present.
115        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
116
117    Returns:
118        The segmentation dataset.
119    """
120    raw_paths, label_paths = get_totalsegmentator_mri_paths(path, split, download)
121
122    if resize_inputs:
123        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
124        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
125            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
126        )
127
128    return torch_em.default_segmentation_dataset(
129        raw_paths=raw_paths,
130        raw_key="data",
131        label_paths=label_paths,
132        label_key="data",
133        patch_shape=patch_shape,
134        is_seg_dataset=True,
135        **kwargs
136    )
137
138
139def get_totalsegmentator_mri_loader(
140    path: Union[os.PathLike, str],
141    batch_size: int,
142    patch_shape: Tuple[int, ...],
143    split: Literal['train', 'test'],
144    resize_inputs: bool = False,
145    download: bool = False,
146    **kwargs
147) -> DataLoader:
148    """Get the TotalSegmentator MRI dataloader for segmentation of anatomical structures in MRI.
149
150    Args:
151        path: Filepath to a folder where the data is downloaded for further processing.
152        batch_size: The batch size for training.
153        patch_shape: The patch shape to use for training.
154        split: The choice of data split. Either 'train' or 'test'.
155        resize_inputs: Whether to resize inputs to the desired patch shape.
156        download: Whether to download the data if it is not present.
157        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
158
159    Returns:
160        The DataLoader.
161    """
162    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
163    dataset = get_totalsegmentator_mri_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
164    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)