torch_em.data.datasets.medical.atm22

  1"""The ATM dataset contains annotations for airway tree segmentation in chest CT scans.
  2
  3It comprises the training set of the ATM'22 challenge (https://atm22.grand-challenge.org): 300 chest CT
  4volumes with a full airway tree annotation, distributed in two batches on Zenodo under the CC BY 4.0 license.
  5The scans are collected from multiple sites and a part of them stems from LIDC-IDRI and EXACT'09.
  6The images of the 20 EXACT'09 cases (ATM_242 - ATM_250 and ATM_501 - ATM_511) are not redistributed by the
  7organizers, so only the 280 cases with both an image and an annotation are exposed by this dataset.
  8
  9NOTE: The label legend is as follows:
 10- background: 0, airway: 1
 11
 12The dataset is located at https://doi.org/10.5281/zenodo.7949582 (TrainBatch1) and
 13https://doi.org/10.5281/zenodo.7949571 (TrainBatch2). See https://atm22.grand-challenge.org for the challenge.
 14
 15This dataset is from the publication https://doi.org/10.1016/j.media.2023.102957.
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20from glob import glob
 21from natsort import natsorted
 22from typing import Union, Tuple, List
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URLS = {
 32    "TrainBatch1": "https://zenodo.org/api/records/7949582/files/TrainBatch1.rar/content",
 33    "TrainBatch2": "https://zenodo.org/api/records/7949571/files/TrainBatch2.rar/content",
 34}
 35
 36CHECKSUMS = {
 37    "TrainBatch1": "355f8c5b19f2b481a2294f0b7c06107893bba198b1ac346d033085eb1d705a20",
 38    "TrainBatch2": "83a2fee661f80c29783811e2dc7913b45aad38ed19f3341c1686a6609fe815e5",
 39}
 40
 41LABEL_IDS = {"background": 0, "airway": 1}
 42
 43
 44def get_atm22_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 45    """Download the ATM dataset.
 46
 47    Args:
 48        path: Filepath to a folder where the data is downloaded for further processing.
 49        download: Whether to download the data if it is not present.
 50
 51    Returns:
 52        Filepath where the data is stored.
 53    """
 54    os.makedirs(path, exist_ok=True)
 55
 56    for batch, url in URLS.items():
 57        if os.path.exists(os.path.join(path, batch)):
 58            continue
 59
 60        rar_path = os.path.join(path, f"{batch}.rar")
 61        util.download_source(path=rar_path, url=url, download=download, checksum=CHECKSUMS[batch])
 62        util.unzip_rarfile(rar_path=rar_path, dst=path, remove=True)
 63
 64    return path
 65
 66
 67def get_atm22_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 68    """Get paths to the ATM data.
 69
 70    Args:
 71        path: Filepath to a folder where the data is downloaded for further processing.
 72        download: Whether to download the data if it is not present.
 73
 74    Returns:
 75        List of filepaths for the image data.
 76        List of filepaths for the label data.
 77    """
 78    data_dir = get_atm22_data(path, download)
 79
 80    raw_paths = natsorted(glob(os.path.join(data_dir, "TrainBatch*", "imagesTr", "*.nii.gz")))
 81    label_paths = [p.replace(os.sep + "imagesTr" + os.sep, os.sep + "labelsTr" + os.sep) for p in raw_paths]
 82    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 83
 84    return raw_paths, label_paths
 85
 86
 87def get_atm22_dataset(
 88    path: Union[os.PathLike, str],
 89    patch_shape: Tuple[int, ...],
 90    resize_inputs: bool = False,
 91    download: bool = False,
 92    **kwargs
 93) -> Dataset:
 94    """Get the ATM dataset for airway tree segmentation.
 95
 96    Args:
 97        path: Filepath to a folder where the data is downloaded for further processing.
 98        patch_shape: The patch shape to use for training.
 99        resize_inputs: Whether to resize inputs to the desired patch shape.
100        download: Whether to download the data if it is not present.
101        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
102
103    Returns:
104        The segmentation dataset.
105    """
106    raw_paths, label_paths = get_atm22_paths(path, download)
107
108    if resize_inputs:
109        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
110        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
111            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
112        )
113
114    return torch_em.default_segmentation_dataset(
115        raw_paths=raw_paths,
116        raw_key="data",
117        label_paths=label_paths,
118        label_key="data",
119        patch_shape=patch_shape,
120        is_seg_dataset=True,
121        **kwargs
122    )
123
124
125def get_atm22_loader(
126    path: Union[os.PathLike, str],
127    batch_size: int,
128    patch_shape: Tuple[int, ...],
129    resize_inputs: bool = False,
130    download: bool = False,
131    **kwargs
132) -> DataLoader:
133    """Get the ATM dataloader for airway tree segmentation.
134
135    Args:
136        path: Filepath to a folder where the data is downloaded for further processing.
137        batch_size: The batch size for training.
138        patch_shape: The patch shape to use for training.
139        resize_inputs: Whether to resize inputs to the desired patch shape.
140        download: Whether to download the data if it is not present.
141        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
142
143    Returns:
144        The DataLoader.
145    """
146    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
147    dataset = get_atm22_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
148    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)