torch_em.data.datasets.medical.totalsegmentator

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