torch_em.data.datasets.medical.aeropath

  1"""The AeroPath dataset contains annotations for airway and lung segmentation in thoracic CT
  2(computed tomography angiography) scans of lung cancer patients with challenging pathologies.
  3
  4The dataset comprises 27 CT volumes with binary annotations for the airways and for the lungs.
  5
  6NOTE: The label legend is as follows:
  7- background: 0, airways: 1 (for 'label_choice' = 'airways')
  8- background: 0, lungs: 1 (for 'label_choice' = 'lungs')
  9
 10The dataset is located at https://zenodo.org/records/10069289 (see also https://github.com/raidionics/AeroPath).
 11
 12This dataset is from the publication https://doi.org/10.48550/arXiv.2311.01138.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, Literal, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://zenodo.org/records/10069289/files/AeroPath.zip?download=1"
 29CHECKSUM = "996b6bd7c79b71a871568293bf6927a52e8e73c1a4dadc1b0975ce3eec3e42ee"
 30
 31
 32def get_aeropath_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the AeroPath 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 where the data is stored.
 41    """
 42    data_dir = os.path.join(path, "AeroPath")
 43    if os.path.exists(data_dir):
 44        return data_dir
 45
 46    os.makedirs(path, exist_ok=True)
 47
 48    zip_path = os.path.join(path, "AeroPath.zip")
 49    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 50    util.unzip(zip_path=zip_path, dst=path)
 51
 52    return data_dir
 53
 54
 55def get_aeropath_paths(
 56    path: Union[os.PathLike, str], label_choice: Literal["airways", "lungs"] = "airways", download: bool = False
 57) -> Tuple[List[str], List[str]]:
 58    """Get paths to the AeroPath data.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        label_choice: The choice of annotated structure. Either 'airways' or 'lungs'.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    if label_choice not in ["airways", "lungs"]:
 70        raise ValueError(f"'{label_choice}' is not a valid label choice. Please choose from 'airways' or 'lungs'.")
 71
 72    data_dir = get_aeropath_data(path, download)
 73
 74    raw_paths = natsorted(glob(os.path.join(data_dir, "*", "*_CT_HR.nii.gz")))
 75    label_paths = [p.replace("_CT_HR.nii.gz", f"_CT_HR_label_{label_choice}.nii.gz") for p in raw_paths]
 76    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 77
 78    return raw_paths, label_paths
 79
 80
 81def get_aeropath_dataset(
 82    path: Union[os.PathLike, str],
 83    patch_shape: Tuple[int, ...],
 84    label_choice: Literal["airways", "lungs"] = "airways",
 85    resize_inputs: bool = False,
 86    download: bool = False,
 87    **kwargs
 88) -> Dataset:
 89    """Get the AeroPath dataset for airway (or lung) segmentation.
 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        label_choice: The choice of annotated structure. Either 'airways' or 'lungs'.
 95        resize_inputs: Whether to resize inputs to the desired patch shape.
 96        download: Whether to download the data if it is not present.
 97        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 98
 99    Returns:
100        The segmentation dataset.
101    """
102    raw_paths, label_paths = get_aeropath_paths(path, label_choice, download)
103
104    if resize_inputs:
105        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
106        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
107            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
108        )
109
110    return torch_em.default_segmentation_dataset(
111        raw_paths=raw_paths,
112        raw_key="data",
113        label_paths=label_paths,
114        label_key="data",
115        patch_shape=patch_shape,
116        is_seg_dataset=True,
117        **kwargs
118    )
119
120
121def get_aeropath_loader(
122    path: Union[os.PathLike, str],
123    batch_size: int,
124    patch_shape: Tuple[int, ...],
125    label_choice: Literal["airways", "lungs"] = "airways",
126    resize_inputs: bool = False,
127    download: bool = False,
128    **kwargs
129) -> DataLoader:
130    """Get the AeroPath dataloader for airway (or lung) 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        label_choice: The choice of annotated structure. Either 'airways' or 'lungs'.
137        resize_inputs: Whether to resize inputs to the desired patch shape.
138        download: Whether to download the data if it is not present.
139        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
140
141    Returns:
142        The DataLoader.
143    """
144    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
145    dataset = get_aeropath_dataset(path, patch_shape, label_choice, resize_inputs, download, **ds_kwargs)
146    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)