torch_em.data.datasets.medical.totalsegmentator_liver_lesions

  1"""The TotalSegmentator liver lesions dataset contains annotations for focal liver lesions in CT scans.
  2
  3This is the training dataset for the "liver_lesions" task of the TotalSegmentator repository
  4(https://github.com/wasserth/TotalSegmentator), which is distributed separately from the main
  5TotalSegmentator dataset (see `torch_em.data.datasets.medical.totalsegmentator`). It consists of 842
  6CT 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.20272572 and licensed under CC BY 4.0.
  9
 10This dataset is part of the TotalSegmentator project, published at https://doi.org/10.1148/ryai.230024.
 11Please 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/20272572/files/Dataset591_liver_lesions.zip"
 26CHECKSUM = "b9cc1132407f003a6c98b9cbdfab9fff904e995d12fd27b9f6b83dbdde0f5a3a"
 27
 28LABEL_IDS = {"background": 0, "liver_lesions": 1}
 29
 30
 31def get_totalsegmentator_liver_lesions_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the TotalSegmentator liver lesions 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, "Dataset591_liver_lesions.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_paths(
 55    path: Union[os.PathLike, str], download: bool = False
 56) -> Tuple[List[str], List[str]]:
 57    """Get paths to the TotalSegmentator liver lesions 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    data_dir = get_totalsegmentator_liver_lesions_data(path, download)
 68
 69    raw_paths, label_paths = [], []
 70    for raw_path in sorted(glob(os.path.join(data_dir, "imagesTr", "*_0000.nii.gz"))):
 71        case_id = os.path.basename(raw_path)[:-len("_0000.nii.gz")]
 72        label_path = os.path.join(data_dir, "labelsTr", f"{case_id}.nii.gz")
 73        assert os.path.exists(label_path), label_path
 74        raw_paths.append(raw_path)
 75        label_paths.append(label_path)
 76
 77    assert len(raw_paths) > 0
 78    return raw_paths, label_paths
 79
 80
 81def get_totalsegmentator_liver_lesions_dataset(
 82    path: Union[os.PathLike, str],
 83    patch_shape: Tuple[int, ...],
 84    resize_inputs: bool = False,
 85    download: bool = False,
 86    **kwargs
 87) -> Dataset:
 88    """Get the TotalSegmentator liver lesions dataset for focal liver lesion segmentation in CT.
 89
 90    Args:
 91        path: Filepath to a folder where the data is downloaded for further processing.
 92        patch_shape: The patch shape to use for training.
 93        resize_inputs: Whether to resize inputs to the desired patch shape.
 94        download: Whether to download the data if it is not present.
 95        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 96
 97    Returns:
 98        The segmentation dataset.
 99    """
100    raw_paths, label_paths = get_totalsegmentator_liver_lesions_paths(path, download)
101
102    if resize_inputs:
103        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
104        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
105            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
106        )
107
108    return torch_em.default_segmentation_dataset(
109        raw_paths=raw_paths,
110        raw_key="data",
111        label_paths=label_paths,
112        label_key="data",
113        patch_shape=patch_shape,
114        is_seg_dataset=True,
115        **kwargs
116    )
117
118
119def get_totalsegmentator_liver_lesions_loader(
120    path: Union[os.PathLike, str],
121    batch_size: int,
122    patch_shape: Tuple[int, ...],
123    resize_inputs: bool = False,
124    download: bool = False,
125    **kwargs
126) -> DataLoader:
127    """Get the TotalSegmentator liver lesions dataloader for focal liver lesion segmentation in CT.
128
129    Args:
130        path: Filepath to a folder where the data is downloaded for further processing.
131        batch_size: The batch size for training.
132        patch_shape: The patch shape to use for training.
133        resize_inputs: Whether to resize inputs to the desired patch shape.
134        download: Whether to download the data if it is not present.
135        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
136
137    Returns:
138        The DataLoader.
139    """
140    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
141    dataset = get_totalsegmentator_liver_lesions_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
142    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)