torch_em.data.datasets.medical.segthor

  1"""The SegTHOR dataset contains annotations for thoracic organ-at-risk segmentation in CT scans
  2of lung or esophageal cancer patients.
  3
  4It comprises the training set of the SegTHOR challenge (https://competitions.codalab.org/competitions/21145):
  540 thoracic CT scans with a dense annotation of the esophagus, heart, trachea and aorta. The 20 test volumes
  6of the challenge are distributed without annotations and are therefore not included here.
  7
  8NOTE: The label legend is as follows:
  9- background: 0, esophagus: 1, heart: 2, trachea: 3, aorta: 4
 10The ids were verified on the data: label 2 is by far the largest structure and sits anteriorly in the lower
 11thorax (heart), label 3 is confined to the upper thorax anterior to label 1 (trachea vs. esophagus), and
 12label 4 spans the whole craniocaudal extent (aorta).
 13
 14The data is a redistribution of the official challenge data at https://doi.org/10.5281/zenodo.16663661
 15(CC BY 4.0). The official release at https://competitions.codalab.org/competitions/21145 requires a signed
 16user agreement, so please make sure that you are allowed to use the data for your purpose.
 17
 18This dataset is from the publication https://doi.org/10.1109/ICPR48806.2021.9411873.
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23from glob import glob
 24from natsort import natsorted
 25from typing import Union, Tuple, List
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URL = "https://zenodo.org/records/16663661/files/SegTHOR.zip"
 35CHECKSUM = "5dba9bcc681c895b7d7b59cf5a26cae570bc8bba02c74f9cb128f3292410fb15"
 36
 37LABEL_IDS = {"background": 0, "esophagus": 1, "heart": 2, "trachea": 3, "aorta": 4}
 38
 39
 40def get_segthor_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 41    """Download the SegTHOR dataset.
 42
 43    Args:
 44        path: Filepath to a folder where the data is downloaded for further processing.
 45        download: Whether to download the data if it is not present.
 46
 47    Returns:
 48        Filepath where the data is stored.
 49    """
 50    data_dir = os.path.join(path, "SegTHOR")
 51    if os.path.exists(data_dir):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55
 56    zip_path = os.path.join(path, "SegTHOR.zip")
 57    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 58    util.unzip(zip_path=zip_path, dst=path)
 59
 60    return data_dir
 61
 62
 63def get_segthor_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 64    """Get paths to the SegTHOR data.
 65
 66    Args:
 67        path: Filepath to a folder where the data is downloaded for further processing.
 68        download: Whether to download the data if it is not present.
 69
 70    Returns:
 71        List of filepaths for the image data.
 72        List of filepaths for the label data.
 73    """
 74    data_dir = get_segthor_data(path, download)
 75
 76    case_dirs = natsorted(glob(os.path.join(data_dir, "Patient_*")))
 77    raw_paths = [os.path.join(p, f"{os.path.basename(p)}.nii.gz") for p in case_dirs]
 78    label_paths = [os.path.join(p, "GT.nii.gz") for p in case_dirs]
 79    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths + label_paths)
 80
 81    return raw_paths, label_paths
 82
 83
 84def get_segthor_dataset(
 85    path: Union[os.PathLike, str],
 86    patch_shape: Tuple[int, ...],
 87    resize_inputs: bool = False,
 88    download: bool = False,
 89    **kwargs
 90) -> Dataset:
 91    """Get the SegTHOR dataset for thoracic organ-at-risk segmentation.
 92
 93    Args:
 94        path: Filepath to a folder where the data is downloaded for further processing.
 95        patch_shape: The patch shape to use for training.
 96        resize_inputs: Whether to resize inputs to the desired patch shape.
 97        download: Whether to download the data if it is not present.
 98        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 99
100    Returns:
101        The segmentation dataset.
102    """
103    raw_paths, label_paths = get_segthor_paths(path, download)
104
105    if resize_inputs:
106        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
107        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
108            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
109        )
110
111    return torch_em.default_segmentation_dataset(
112        raw_paths=raw_paths,
113        raw_key="data",
114        label_paths=label_paths,
115        label_key="data",
116        patch_shape=patch_shape,
117        is_seg_dataset=True,
118        **kwargs
119    )
120
121
122def get_segthor_loader(
123    path: Union[os.PathLike, str],
124    batch_size: int,
125    patch_shape: Tuple[int, ...],
126    resize_inputs: bool = False,
127    download: bool = False,
128    **kwargs
129) -> DataLoader:
130    """Get the SegTHOR dataloader for thoracic organ-at-risk segmentation.
131
132    Args:
133        path: Filepath to a folder where the data is downloaded for further processing.
134        batch_size: The batch size for training.
135        patch_shape: The patch shape to use for training.
136        resize_inputs: Whether to resize inputs to the desired patch shape.
137        download: Whether to download the data if it is not present.
138        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
139
140    Returns:
141        The DataLoader.
142    """
143    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
144    dataset = get_segthor_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
145    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)