torch_em.data.datasets.medical.totalsegmentator_effusion

  1"""The TotalSegmentator effusion dataset contains annotations for pleural and pericardial effusion in CT scans.
  2
  3This is the training dataset for the "pleural_pericard_effusion" 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 600
  6CT volumes with a semantic label volume per case, following the class order of the "pleural_pericard_effusion"
  7task: 0 = background, 1 = lung_pleural, 2 = pleural_effusion, 3 = pericardial_effusion (see `LABEL_IDS`).
  8
  9The dataset is located at https://doi.org/10.5281/zenodo.20272295 and licensed under CC BY 4.0.
 10
 11This dataset is part of the TotalSegmentator project, published at https://doi.org/10.1148/ryai.230024.
 12Please cite it if you use this dataset in your research.
 13"""
 14
 15import os
 16from glob import glob
 17from typing import Union, Tuple, List
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26URL = "https://zenodo.org/records/20272295/files/Dataset315_pleural_pericard_effusion.zip"
 27CHECKSUM = "1c8f043b3e466b49af80fa87df6b048b45dcd675796e79a8f3cc2abc18f3bba7"
 28
 29LABEL_IDS = {"background": 0, "lung_pleural": 1, "pleural_effusion": 2, "pericardial_effusion": 3}
 30
 31
 32def get_totalsegmentator_effusion_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the TotalSegmentator pleural / pericardial effusion dataset.
 34
 35    Args:
 36        path: Filepath to a folder where the data is downloaded for further processing.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath to the folder with the 'imagesTr' and 'labelsTr' folders.
 41    """
 42    # The archive has no top-level folder, hence it is extracted directly into 'path'.
 43    data_dir = path
 44    if os.path.exists(os.path.join(data_dir, "dataset.json")):
 45        return data_dir
 46
 47    os.makedirs(path, exist_ok=True)
 48    zip_path = os.path.join(path, "Dataset315_pleural_pericard_effusion.zip")
 49    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 50    util.unzip(zip_path=zip_path, dst=data_dir)
 51
 52    return data_dir
 53
 54
 55def get_totalsegmentator_effusion_paths(
 56    path: Union[os.PathLike, str], download: bool = False
 57) -> Tuple[List[str], List[str]]:
 58    """Get paths to the TotalSegmentator effusion data.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        download: Whether to download the data if it is not present.
 63
 64    Returns:
 65        List of filepaths for the image data.
 66        List of filepaths for the label data.
 67    """
 68    data_dir = get_totalsegmentator_effusion_data(path, download)
 69
 70    raw_paths, label_paths = [], []
 71    for raw_path in sorted(glob(os.path.join(data_dir, "imagesTr", "*_0000.nii.gz"))):
 72        case_id = os.path.basename(raw_path)[:-len("_0000.nii.gz")]
 73        label_path = os.path.join(data_dir, "labelsTr", f"{case_id}.nii.gz")
 74        assert os.path.exists(label_path), label_path
 75        raw_paths.append(raw_path)
 76        label_paths.append(label_path)
 77
 78    assert len(raw_paths) > 0
 79    return raw_paths, label_paths
 80
 81
 82def get_totalsegmentator_effusion_dataset(
 83    path: Union[os.PathLike, str],
 84    patch_shape: Tuple[int, ...],
 85    resize_inputs: bool = False,
 86    download: bool = False,
 87    **kwargs
 88) -> Dataset:
 89    """Get the TotalSegmentator effusion dataset for pleural / pericardial effusion segmentation in CT.
 90
 91    Args:
 92        path: Filepath to a folder where the data is downloaded for further processing.
 93        patch_shape: The patch shape to use for training.
 94        resize_inputs: Whether to resize inputs to the desired patch shape.
 95        download: Whether to download the data if it is not present.
 96        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 97
 98    Returns:
 99        The segmentation dataset.
100    """
101    raw_paths, label_paths = get_totalsegmentator_effusion_paths(path, download)
102
103    if resize_inputs:
104        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
105        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
106            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
107        )
108
109    return torch_em.default_segmentation_dataset(
110        raw_paths=raw_paths,
111        raw_key="data",
112        label_paths=label_paths,
113        label_key="data",
114        patch_shape=patch_shape,
115        is_seg_dataset=True,
116        **kwargs
117    )
118
119
120def get_totalsegmentator_effusion_loader(
121    path: Union[os.PathLike, str],
122    batch_size: int,
123    patch_shape: Tuple[int, ...],
124    resize_inputs: bool = False,
125    download: bool = False,
126    **kwargs
127) -> DataLoader:
128    """Get the TotalSegmentator effusion dataloader for pleural / pericardial effusion segmentation in CT.
129
130    Args:
131        path: Filepath to a folder where the data is downloaded for further processing.
132        batch_size: The batch size for training.
133        patch_shape: The patch shape to use for training.
134        resize_inputs: Whether to resize inputs to the desired patch shape.
135        download: Whether to download the data if it is not present.
136        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
137
138    Returns:
139        The DataLoader.
140    """
141    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
142    dataset = get_totalsegmentator_effusion_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
143    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)