torch_em.data.datasets.medical.totalsegmentator_liver_segments_mr

  1"""The TotalSegmentator liver segments MRI dataset contains annotations for the 8 Couinaud liver segments
  2in MRI scans.
  3
  4This is the training dataset for the "liver_segments_mr" task of the TotalSegmentator repository
  5(https://github.com/wasserth/TotalSegmentator), which is distributed separately from the main
  6TotalSegmentator MRI dataset (see `torch_em.data.datasets.medical.totalsegmentator_mri`). It consists of
  7115 MRI volumes, of which about half are negative controls with an entirely empty (all-background) label
  8volume; `get_totalsegmentator_liver_segments_mr_paths` filters these out, so only volumes with manual
  9segmentations of the 8 Couinaud liver segments are returned.
 10
 11The dataset is located at https://doi.org/10.5281/zenodo.20746820 and licensed under CC BY 4.0.
 12
 13This dataset is part of the TotalSegmentator project, published at
 14https://doi.org/10.1007/s10278-025-01716-y. Please cite it if you use this dataset in your research.
 15"""
 16
 17import os
 18from glob import glob
 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://zenodo.org/records/20746820/files/Dataset576_liver_segments_mr.zip"
 29CHECKSUM = "621dc9cfe51348980cf1ef65eaf662cba2b18c499d27b186a236429aa954e438"
 30
 31LABEL_IDS = {
 32    "background": 0,
 33    "liver_segment_1": 1,
 34    "liver_segment_2": 2,
 35    "liver_segment_3": 3,
 36    "liver_segment_4": 4,
 37    "liver_segment_5": 5,
 38    "liver_segment_6": 6,
 39    "liver_segment_7": 7,
 40    "liver_segment_8": 8,
 41}
 42
 43
 44def get_totalsegmentator_liver_segments_mr_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 45    """Download the TotalSegmentator liver segments MRI dataset.
 46
 47    Args:
 48        path: Filepath to a folder where the data is downloaded for further processing.
 49        download: Whether to download the data if it is not present.
 50
 51    Returns:
 52        Filepath to the folder with the 'imagesTr' and 'labelsTr' folders.
 53    """
 54    data_dir = os.path.join(path, "Dataset576_liver_segments_mr")
 55    if os.path.exists(os.path.join(data_dir, "dataset.json")):
 56        return data_dir
 57
 58    os.makedirs(path, exist_ok=True)
 59    zip_path = os.path.join(path, "Dataset576_liver_segments_mr.zip")
 60    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 61    util.unzip(zip_path=zip_path, dst=path)
 62
 63    return data_dir
 64
 65
 66def get_totalsegmentator_liver_segments_mr_paths(
 67    path: Union[os.PathLike, str], download: bool = False
 68) -> Tuple[List[str], List[str]]:
 69    """Get paths to the TotalSegmentator liver segments MRI data.
 70
 71    Args:
 72        path: Filepath to a folder where the data is downloaded for further processing.
 73        download: Whether to download the data if it is not present.
 74
 75    Returns:
 76        List of filepaths for the image data.
 77        List of filepaths for the label data.
 78    """
 79    import nibabel as nib
 80    import numpy as np
 81
 82    data_dir = get_totalsegmentator_liver_segments_mr_data(path, download)
 83
 84    raw_paths, label_paths = [], []
 85    for raw_path in sorted(glob(os.path.join(data_dir, "imagesTr", "*_0000.nii.gz"))):
 86        case_id = os.path.basename(raw_path)[:-len("_0000.nii.gz")]
 87        label_path = os.path.join(data_dir, "labelsTr", f"{case_id}.nii.gz")
 88        assert os.path.exists(label_path), label_path
 89
 90        # Skip the negative control cases, whose label volume is entirely background.
 91        if not np.any(nib.load(label_path).get_fdata()):
 92            continue
 93
 94        raw_paths.append(raw_path)
 95        label_paths.append(label_path)
 96
 97    assert len(raw_paths) > 0
 98    return raw_paths, label_paths
 99
100
101def get_totalsegmentator_liver_segments_mr_dataset(
102    path: Union[os.PathLike, str],
103    patch_shape: Tuple[int, ...],
104    resize_inputs: bool = False,
105    download: bool = False,
106    **kwargs
107) -> Dataset:
108    """Get the TotalSegmentator liver segments MRI dataset for Couinaud liver segment segmentation in MRI.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        patch_shape: The patch shape to use for training.
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_liver_segments_mr_paths(path, 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_liver_segments_mr_loader(
140    path: Union[os.PathLike, str],
141    batch_size: int,
142    patch_shape: Tuple[int, ...],
143    resize_inputs: bool = False,
144    download: bool = False,
145    **kwargs
146) -> DataLoader:
147    """Get the TotalSegmentator liver segments MRI dataloader for Couinaud liver segment segmentation in MRI.
148
149    Args:
150        path: Filepath to a folder where the data is downloaded for further processing.
151        batch_size: The batch size for training.
152        patch_shape: The patch shape to use for training.
153        resize_inputs: Whether to resize inputs to the desired patch shape.
154        download: Whether to download the data if it is not present.
155        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
156
157    Returns:
158        The DataLoader.
159    """
160    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
161    dataset = get_totalsegmentator_liver_segments_mr_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
162    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)