torch_em.data.datasets.medical.totalsegmentator_liver_lesions_mr

  1"""The TotalSegmentator liver lesions MRI dataset contains annotations for focal liver lesions in MRI scans.
  2
  3This is the training dataset for the "liver_lesions_mr" task of the TotalSegmentator repository
  4(https://github.com/wasserth/TotalSegmentator), which is distributed separately from the main
  5TotalSegmentator MRI dataset (see `torch_em.data.datasets.medical.totalsegmentator_mri`). It consists of
  6MRI volumes with a single binary label for focal liver lesions (0 = background, 1 = liver lesion).
  7
  8The dataset is located at https://doi.org/10.5281/zenodo.20272348 and licensed under CC BY 4.0.
  9
 10This dataset is part of the TotalSegmentator project, published at
 11https://doi.org/10.1007/s10278-025-01716-y. Please cite it if you use this dataset in your research.
 12"""
 13
 14import os
 15from glob import glob
 16from typing import Union, Tuple, List
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://zenodo.org/records/20272348/files/Dataset589_liver_lesions_mr.zip"
 26CHECKSUM = "4613f47e879a04fcbf2a8792eaff2cd4aa6b324185d69e22100e1c071e2e5b05"
 27
 28LABEL_IDS = {"background": 0, "liver_lesion": 1}
 29
 30
 31def get_totalsegmentator_liver_lesions_mr_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the TotalSegmentator liver lesions MRI dataset.
 33
 34    Args:
 35        path: Filepath to a folder where the data is downloaded for further processing.
 36        download: Whether to download the data if it is not present.
 37
 38    Returns:
 39        Filepath to the folder with the 'imagesTr' and 'labelsTr' folders.
 40    """
 41    # The archive has no top-level folder, hence it is extracted directly into 'path'.
 42    data_dir = path
 43    if os.path.exists(os.path.join(data_dir, "dataset.json")):
 44        return data_dir
 45
 46    os.makedirs(path, exist_ok=True)
 47    zip_path = os.path.join(path, "Dataset589_liver_lesions_mr.zip")
 48    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 49    util.unzip(zip_path=zip_path, dst=data_dir)
 50
 51    return data_dir
 52
 53
 54def get_totalsegmentator_liver_lesions_mr_paths(
 55    path: Union[os.PathLike, str], download: bool = False
 56) -> Tuple[List[str], List[str]]:
 57    """Get paths to the TotalSegmentator liver lesions MRI data.
 58
 59    Args:
 60        path: Filepath to a folder where the data is downloaded for further processing.
 61        download: Whether to download the data if it is not present.
 62
 63    Returns:
 64        List of filepaths for the image data.
 65        List of filepaths for the label data.
 66    """
 67    import nibabel as nib
 68    import numpy as np
 69
 70    data_dir = get_totalsegmentator_liver_lesions_mr_data(path, download)
 71
 72    raw_paths, label_paths = [], []
 73    for raw_path in sorted(glob(os.path.join(data_dir, "imagesTr", "*_0000.nii.gz"))):
 74        case_id = os.path.basename(raw_path)[:-len("_0000.nii.gz")]
 75        label_path = os.path.join(data_dir, "labelsTr", f"{case_id}.nii.gz")
 76        assert os.path.exists(label_path), label_path
 77
 78        # Skip the negative control cases, whose label volume is entirely background.
 79        if not np.any(nib.load(label_path).get_fdata()):
 80            continue
 81
 82        raw_paths.append(raw_path)
 83        label_paths.append(label_path)
 84
 85    assert len(raw_paths) > 0
 86    return raw_paths, label_paths
 87
 88
 89def get_totalsegmentator_liver_lesions_mr_dataset(
 90    path: Union[os.PathLike, str],
 91    patch_shape: Tuple[int, ...],
 92    resize_inputs: bool = False,
 93    download: bool = False,
 94    **kwargs
 95) -> Dataset:
 96    """Get the TotalSegmentator liver lesions MRI dataset for focal liver lesion segmentation in MRI.
 97
 98    Args:
 99        path: Filepath to a folder where the data is downloaded for further processing.
100        patch_shape: The patch shape to use for training.
101        resize_inputs: Whether to resize inputs to the desired patch shape.
102        download: Whether to download the data if it is not present.
103        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
104
105    Returns:
106        The segmentation dataset.
107    """
108    raw_paths, label_paths = get_totalsegmentator_liver_lesions_mr_paths(path, download)
109
110    if resize_inputs:
111        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
112        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
113            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
114        )
115
116    return torch_em.default_segmentation_dataset(
117        raw_paths=raw_paths,
118        raw_key="data",
119        label_paths=label_paths,
120        label_key="data",
121        patch_shape=patch_shape,
122        is_seg_dataset=True,
123        **kwargs
124    )
125
126
127def get_totalsegmentator_liver_lesions_mr_loader(
128    path: Union[os.PathLike, str],
129    batch_size: int,
130    patch_shape: Tuple[int, ...],
131    resize_inputs: bool = False,
132    download: bool = False,
133    **kwargs
134) -> DataLoader:
135    """Get the TotalSegmentator liver lesions MRI dataloader for focal liver lesion segmentation in MRI.
136
137    Args:
138        path: Filepath to a folder where the data is downloaded for further processing.
139        batch_size: The batch size for training.
140        patch_shape: The patch shape to use for training.
141        resize_inputs: Whether to resize inputs to the desired patch shape.
142        download: Whether to download the data if it is not present.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
144
145    Returns:
146        The DataLoader.
147    """
148    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
149    dataset = get_totalsegmentator_liver_lesions_mr_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
150    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)