torch_em.data.datasets.medical.totalsegmentator

  1"""The TotalSegmentator dataset contains annotations for 117 anatomical structures in CT scans.
  2
  3Two versions of the dataset can be selected via the 'version' argument:
  4- 'v2' (default, for backward compatibility): 1228 CT volumes, located at https://doi.org/10.5281/zenodo.10047292.
  5- 'v3': 1939 CT volumes (291 additional pediatric CTs on top of the 1228 from 'v2', plus corrected and refined
  6  labels, especially for bones and vertebrae), located at https://doi.org/10.5281/zenodo.22688904.
  7
  8Both versions ship an official split in 'meta.csv' ('v2': train / val / test, 'v3': train / test only, i.e. no
  9'val' split is available for 'v3') and provide each anatomical structure as a separate binary mask.
 10`get_totalsegmentator_data` merges these masks into a single semantic label volume per case, where the label id
 11of each structure is its (1-based) position in `CLASS_NAMES` for 'v2', or `CLASS_NAMES_V3` for 'v3' (see
 12`CLASS_IDS` for the 'v2' name -> id mapping). This is the class order of the 'total' task in the
 13TotalSegmentator repository (https://github.com/wasserth/TotalSegmentator). 'v3' relabels 'vertebrae_S1' as
 14'vertebrae_L6' as part of its corrected vertebrae labeling (verified on the released 'v3' archive), all other
 15116 structures are unchanged between 'v2' and 'v3'. The masks of a few structures may overlap, in this case the
 16structure with the higher label id takes precedence.
 17
 18This dataset is from the publication https://doi.org/10.1148/ryai.230024.
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23from glob import glob
 24from concurrent import futures
 25from typing import Union, Tuple, Literal, List, Optional
 26
 27import numpy as np
 28from tqdm import tqdm
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URLS = {
 38    "v2": "https://zenodo.org/records/10047292/files/Totalsegmentator_dataset_v201.zip",
 39    "v3": "https://zenodo.org/records/22688904/files/Totalsegmentator_dataset_v300.zip",
 40}
 41CHECKSUMS = {
 42    "v2": "741dbc911a768e2ac2671c66d55332f7302ad624c915a57d08b142d8bdf0ca26",
 43    "v3": "b56ae18553853ff256fb0eef3a02e322fbc6c862d9015238c8c46b1b76fa027b",
 44}
 45DATA_DIRNAMES = {
 46    "v2": "Totalsegmentator_dataset_v201",
 47    "v3": "Totalsegmentator_dataset_v300",
 48}
 49# The 'v2' archive has no top-level folder (files extract directly into the data folder), while the
 50# 'v3' archive already contains a top-level folder matching 'DATA_DIRNAMES["v3"]'.
 51HAS_TOP_LEVEL_DIR = {"v2": False, "v3": True}
 52
 53CLASS_NAMES = [
 54    "spleen", "kidney_right", "kidney_left", "gallbladder", "liver", "stomach", "pancreas", "adrenal_gland_right",
 55    "adrenal_gland_left", "lung_upper_lobe_left", "lung_lower_lobe_left", "lung_upper_lobe_right",
 56    "lung_middle_lobe_right", "lung_lower_lobe_right", "esophagus", "trachea", "thyroid_gland", "small_bowel",
 57    "duodenum", "colon", "urinary_bladder", "prostate", "kidney_cyst_left", "kidney_cyst_right", "sacrum",
 58    "vertebrae_S1", "vertebrae_L5", "vertebrae_L4", "vertebrae_L3", "vertebrae_L2", "vertebrae_L1", "vertebrae_T12",
 59    "vertebrae_T11", "vertebrae_T10", "vertebrae_T9", "vertebrae_T8", "vertebrae_T7", "vertebrae_T6", "vertebrae_T5",
 60    "vertebrae_T4", "vertebrae_T3", "vertebrae_T2", "vertebrae_T1", "vertebrae_C7", "vertebrae_C6", "vertebrae_C5",
 61    "vertebrae_C4", "vertebrae_C3", "vertebrae_C2", "vertebrae_C1", "heart", "aorta", "pulmonary_vein",
 62    "brachiocephalic_trunk", "subclavian_artery_right", "subclavian_artery_left", "common_carotid_artery_right",
 63    "common_carotid_artery_left", "brachiocephalic_vein_left", "brachiocephalic_vein_right", "atrial_appendage_left",
 64    "superior_vena_cava", "inferior_vena_cava", "portal_vein_and_splenic_vein", "iliac_artery_left",
 65    "iliac_artery_right", "iliac_vena_left", "iliac_vena_right", "humerus_left", "humerus_right", "scapula_left",
 66    "scapula_right", "clavicula_left", "clavicula_right", "femur_left", "femur_right", "hip_left", "hip_right",
 67    "spinal_cord", "gluteus_maximus_left", "gluteus_maximus_right", "gluteus_medius_left", "gluteus_medius_right",
 68    "gluteus_minimus_left", "gluteus_minimus_right", "autochthon_left", "autochthon_right", "iliopsoas_left",
 69    "iliopsoas_right", "brain", "skull", "rib_left_1", "rib_left_2", "rib_left_3", "rib_left_4", "rib_left_5",
 70    "rib_left_6", "rib_left_7", "rib_left_8", "rib_left_9", "rib_left_10", "rib_left_11", "rib_left_12", "rib_right_1",
 71    "rib_right_2", "rib_right_3", "rib_right_4", "rib_right_5", "rib_right_6", "rib_right_7", "rib_right_8",
 72    "rib_right_9", "rib_right_10", "rib_right_11", "rib_right_12", "sternum", "costal_cartilages",
 73]
 74"""The anatomical structures of the TotalSegmentator CT dataset. The label id of a structure is its 1-based index."""
 75
 76CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 77"""Mapping from the name of an anatomical structure to its label id in the merged label volumes."""
 78
 79CLASS_NAMES_V3 = [name if name != "vertebrae_S1" else "vertebrae_L6" for name in CLASS_NAMES]
 80"""The anatomical structures of the 'v3' TotalSegmentator CT dataset. Identical to `CLASS_NAMES`, except that
 81'vertebrae_S1' was replaced by 'vertebrae_L6' as part of the corrected vertebrae labeling in 'v3' (verified on
 82the per-case 'segmentations' folders of the released 'v3' archive, all 117 structures are otherwise unchanged).
 83"""
 84
 85CLASS_NAMES_BY_VERSION = {"v2": CLASS_NAMES, "v3": CLASS_NAMES_V3}
 86
 87
 88def merge_segmentations(case_dir: str, class_names: List[str], label_name: str = "labels.nii.gz") -> str:
 89    """Merge the per-class binary masks of one TotalSegmentator case into a single semantic label volume.
 90
 91    The merged volume is stored as nifti next to the image. If it already exists it is not recomputed,
 92    so that a partially finished conversion can be resumed.
 93
 94    Args:
 95        case_dir: The folder of the case, which contains the 'segmentations' sub-folder.
 96        class_names: The class names in label id order (the first class gets id 1).
 97        label_name: The filename of the merged label volume.
 98
 99    Returns:
100        The filepath to the merged label volume.
101    """
102    import nibabel as nib
103
104    label_path = os.path.join(case_dir, label_name)
105    if os.path.exists(label_path):
106        return label_path
107
108    labels, affine, header = None, None, None
109    for class_id, class_name in enumerate(class_names, start=1):
110        mask_nii = nib.load(os.path.join(case_dir, "segmentations", f"{class_name}.nii.gz"))
111        mask = np.asarray(mask_nii.dataobj) > 0
112        if labels is None:
113            labels = np.zeros(mask.shape, dtype="uint8")
114            affine, header = mask_nii.affine, mask_nii.header
115        labels[mask] = class_id
116
117    # Write to a temporary path first, so that an interrupted conversion is not mistaken for a complete one.
118    tmp_path = os.path.join(case_dir, f"{label_name}.incomplete.nii.gz")
119    nib.save(nib.Nifti1Image(labels, affine, header), tmp_path)
120    os.replace(tmp_path, label_path)
121    return label_path
122
123
124def merge_all_segmentations(case_dirs: List[str], class_names: List[str], n_workers: Optional[int] = None) -> None:
125    """Merge the per-class binary masks of all cases into semantic label volumes.
126
127    Args:
128        case_dirs: The case folders to process.
129        class_names: The class names in label id order.
130        n_workers: The number of parallel workers. By default the number of CPUs (at most 16) is used.
131    """
132    if all(os.path.exists(os.path.join(case_dir, "labels.nii.gz")) for case_dir in case_dirs):
133        return
134
135    if n_workers is None:
136        n_cpus = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)
137        n_workers = min(16, n_cpus)
138
139    with futures.ProcessPoolExecutor(n_workers) as pool:
140        tasks = [pool.submit(merge_segmentations, case_dir, class_names) for case_dir in case_dirs]
141        for task in tqdm(futures.as_completed(tasks), total=len(tasks), desc="Merge the per-class segmentations"):
142            task.result()
143
144
145def read_split(meta_csv: str, split: str, valid_splits: Tuple[str, ...] = ("train", "val", "test")) -> List[str]:
146    """Read the case ids of a split from the TotalSegmentator 'meta.csv'.
147
148    Args:
149        meta_csv: The path to the 'meta.csv' file.
150        split: The choice of data split.
151        valid_splits: The splits available in this dataset.
152
153    Returns:
154        The case ids of the split.
155    """
156    import pandas as pd
157
158    if split not in valid_splits:
159        raise ValueError(f"'{split}' is not a valid split. Choose one of {valid_splits}.")
160
161    meta = pd.read_csv(meta_csv, sep=";", encoding="utf-8-sig")
162    return sorted(meta[meta["split"] == split]["image_id"].tolist())
163
164
165def get_totalsegmentator_data(
166    path: Union[os.PathLike, str],
167    download: bool = False,
168    n_workers: Optional[int] = None,
169    version: Literal["v2", "v3"] = "v2",
170) -> str:
171    """Download the TotalSegmentator CT dataset and merge the per-class masks into semantic label volumes.
172
173    Args:
174        path: Filepath to a folder where the data is downloaded for further processing.
175        download: Whether to download the data if it is not present.
176        n_workers: The number of parallel workers for merging the per-class masks.
177        version: The version of the dataset. Either 'v2' (1228 CTs) or 'v3' (1939 CTs, incl. pediatric CTs
178            and refined bone / vertebrae labels).
179
180    Returns:
181        Filepath where the data is downloaded.
182    """
183    if version not in URLS:
184        raise ValueError(f"'{version}' is not a valid version. Choose one of {list(URLS.keys())}.")
185
186    dirname = DATA_DIRNAMES[version]
187    data_dir = os.path.join(path, dirname)
188    if not os.path.exists(os.path.join(data_dir, "meta.csv")):
189        os.makedirs(path, exist_ok=True)
190        zip_path = os.path.join(path, f"{dirname}.zip")
191        util.download_source(path=zip_path, url=URLS[version], download=download, checksum=CHECKSUMS[version])
192        # Extract into 'data_dir' if the archive has no top-level folder, otherwise into 'path'
193        # (the archive's own top-level folder then becomes 'data_dir').
194        util.unzip(zip_path=zip_path, dst=path if HAS_TOP_LEVEL_DIR[version] else data_dir)
195
196    case_dirs = sorted(glob(os.path.join(data_dir, "s*")))
197    merge_all_segmentations(case_dirs, CLASS_NAMES_BY_VERSION[version], n_workers)
198
199    return data_dir
200
201
202def get_totalsegmentator_paths(
203    path: Union[os.PathLike, str],
204    split: Literal['train', 'val', 'test'],
205    download: bool = False,
206    version: Literal["v2", "v3"] = "v2",
207) -> Tuple[List[str], List[str]]:
208    """Get paths to the TotalSegmentator CT data.
209
210    Args:
211        path: Filepath to a folder where the data is downloaded for further processing.
212        split: The choice of data split.
213        download: Whether to download the data if it is not present.
214        version: The version of the dataset. Either 'v2' (1228 CTs) or 'v3' (1939 CTs).
215
216    Returns:
217        List of filepaths for the image data.
218        List of filepaths for the label data.
219    """
220    data_dir = get_totalsegmentator_data(path, download, version=version)
221    case_ids = read_split(os.path.join(data_dir, "meta.csv"), split)
222
223    raw_paths = [os.path.join(data_dir, case_id, "ct.nii.gz") for case_id in case_ids]
224    label_paths = [os.path.join(data_dir, case_id, "labels.nii.gz") for case_id in case_ids]
225    assert all(os.path.exists(p) for p in raw_paths + label_paths)
226
227    return raw_paths, label_paths
228
229
230def get_totalsegmentator_dataset(
231    path: Union[os.PathLike, str],
232    patch_shape: Tuple[int, ...],
233    split: Literal['train', 'val', 'test'],
234    resize_inputs: bool = False,
235    download: bool = False,
236    version: Literal["v2", "v3"] = "v2",
237    **kwargs
238) -> Dataset:
239    """Get the TotalSegmentator dataset for segmentation of anatomical structures in CT.
240
241    Args:
242        path: Filepath to a folder where the data is downloaded for further processing.
243        patch_shape: The patch shape to use for training.
244        split: The choice of data split.
245        resize_inputs: Whether to resize inputs to the desired patch shape.
246        download: Whether to download the data if it is not present.
247        version: The version of the dataset. Either 'v2' (1228 CTs, default) or 'v3' (1939 CTs).
248        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
249
250    Returns:
251        The segmentation dataset.
252    """
253    raw_paths, label_paths = get_totalsegmentator_paths(path, split, download, version=version)
254
255    if resize_inputs:
256        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
257        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
258            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
259        )
260
261    return torch_em.default_segmentation_dataset(
262        raw_paths=raw_paths,
263        raw_key="data",
264        label_paths=label_paths,
265        label_key="data",
266        patch_shape=patch_shape,
267        is_seg_dataset=True,
268        **kwargs
269    )
270
271
272def get_totalsegmentator_loader(
273    path: Union[os.PathLike, str],
274    batch_size: int,
275    patch_shape: Tuple[int, ...],
276    split: Literal['train', 'val', 'test'],
277    resize_inputs: bool = False,
278    download: bool = False,
279    version: Literal["v2", "v3"] = "v2",
280    **kwargs
281) -> DataLoader:
282    """Get the TotalSegmentator dataloader for segmentation of anatomical structures in CT.
283
284    Args:
285        path: Filepath to a folder where the data is downloaded for further processing.
286        batch_size: The batch size for training.
287        patch_shape: The patch shape to use for training.
288        split: The choice of data split.
289        resize_inputs: Whether to resize inputs to the desired patch shape.
290        download: Whether to download the data if it is not present.
291        version: The version of the dataset. Either 'v2' (1228 CTs, default) or 'v3' (1939 CTs).
292        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
293
294    Returns:
295        The DataLoader.
296    """
297    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
298    dataset = get_totalsegmentator_dataset(path, patch_shape, split, resize_inputs, download, version, **ds_kwargs)
299    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)