torch_em.data.datasets.medical.dap_atlas

  1"""The DAP Atlas dataset contains annotations for 142 anatomical structures in whole-body CT scans.
  2
  3The dataset consists of automatically generated (and post-processed with anatomical guidelines) dense anatomical
  4label volumes for 533 CT scans of the AutoPET dataset. The label volumes are distributed via the DAP Atlas repository,
  5the CT scans are the (original resolution) 'CT.nii.gz' volumes of the AutoPET dataset, which is downloaded via
  6`torch_em.data.datasets.medical.autopet`. The names of the label volumes ('AutoPET_<subject id>_<last 5 digits of
  7the study uid>') uniquely identify the matching AutoPET CT.
  8NOTE: The AutoPET download is ~300 GB, but it is shared with the AutoPET dataset (see `autopet_path`).
  9
 10The label id of each structure is given in `CLASS_IDS` (see also the table in the appendix of the publication).
 11The label ids are: 1: left to annotate (unlabeled placeholder), 2: muscles, 3: fat, 4: abdominal tissue,
 125: mediastinal tissue, 6: esophagus, 7: stomach, 8: small bowel, 9: duodenum, 10: colon, 11: rectum, 12: gallbladder,
 1313: liver, 14: pancreas, 15: kidney left, 16: kidney right, 17: bladder, 18: gonads, 19: prostate, 20: uterocervix,
 1421: uterus, 22: breast left, 23: breast right, 24: spinal canal, 25: brain, 26: spleen, 27: adrenal gland left,
 1528: adrenal gland right, 29: thyroid left, 30: thyroid right, 31: thymus, 32-37: gluteus maximus / medius / minimus
 16(left, right), 38-39: iliopsoas (left, right), 40-41: autochthon (left, right), 42: skin, 43-66: vertebrae C1-L5,
 1767-90: costa 1-12 (left, right), 91: rib cartilage, 92: sternum corpus, 93-94: clavicula (left, right),
 1895-96: scapula (left, right), 97-98: humerus (left, right), 99: skull, 100-101: hip (left, right), 102: sacrum,
 19103-104: femur (left, right), 105: heart, 106: heart atrium left, 107: heart tissue, 108: heart atrium right,
 20109: heart myocardium, 110: heart ventricle left, 111: heart ventricle right, 112-113: iliac artery (left, right),
 21114: aorta, 115-116: iliac vena (left, right), 117: inferior vena cava, 118: portal vein and splenic vein,
 22119: celiac trunk, 120-124: lung lobes, 125: bronchus, 126: trachea, 127: pulmonary artery,
 23128-129: cheek (left, right),
 24130-131: eyeball (left, right), 132: nasal cavity, 133-134: common carotid artery (right, left),
 25135: sternum manubrium, 136-137: internal carotid artery (right, left), 138-139: internal jugular vein (right, left),
 26140: brachiocephalic artery, 141-142: brachiocephalic vein (right, left), 143-144: subclavian artery (right, left).
 27
 28The dataset is located at https://github.com/alexanderjaus/AtlasDataset.
 29
 30This dataset is from the publication https://doi.org/10.48550/arXiv.2307.13375.
 31Please cite it (and the AutoPET publication https://doi.org/10.1038/s41597-022-01718-3) if you use this dataset
 32in your research.
 33"""
 34
 35import os
 36from glob import glob
 37from warnings import warn
 38from natsort import natsorted
 39from typing import Union, Tuple, List, Optional
 40
 41from torch.utils.data import Dataset, DataLoader
 42
 43import torch_em
 44
 45from .. import util
 46from .autopet import get_autopet_data
 47
 48
 49URL = "https://drive.google.com/uc?id=1ex0a9eQULLvKPDwijmijX2h49A-ockNy"
 50CHECKSUM = "eea8b7bf2378a5085bc4ef7014bb4d764c9e276a6d8c8a55add19e9bb2b80622"
 51
 52CLASS_NAMES = [
 53    "left_to_annotate", "muscles", "fat", "abdominal_tissue", "mediastinal_tissue", "esophagus", "stomach",
 54    "small_bowel", "duodenum", "colon", "rectum", "gallbladder", "liver", "pancreas", "kidney_left", "kidney_right",
 55    "bladder", "gonads", "prostate", "uterocervix", "uterus", "breast_left", "breast_right", "spinal_canal", "brain",
 56    "spleen", "adrenal_gland_left", "adrenal_gland_right", "thyroid_left", "thyroid_right", "thymus",
 57    "gluteus_maximus_left", "gluteus_maximus_right", "gluteus_medius_left", "gluteus_medius_right",
 58    "gluteus_minimus_left", "gluteus_minimus_right", "iliopsoas_left", "iliopsoas_right", "autochthon_left",
 59    "autochthon_right", "skin", "vertebrae_C1", "vertebrae_C2", "vertebrae_C3", "vertebrae_C4", "vertebrae_C5",
 60    "vertebrae_C6", "vertebrae_C7", "vertebrae_T1", "vertebrae_T2", "vertebrae_T3", "vertebrae_T4", "vertebrae_T5",
 61    "vertebrae_T6", "vertebrae_T7", "vertebrae_T8", "vertebrae_T9", "vertebrae_T10", "vertebrae_T11", "vertebrae_T12",
 62    "vertebrae_L1", "vertebrae_L2", "vertebrae_L3", "vertebrae_L4", "vertebrae_L5", "costa_1_left", "costa_1_right",
 63    "costa_2_left", "costa_2_right", "costa_3_left", "costa_3_right", "costa_4_left", "costa_4_right", "costa_5_left",
 64    "costa_5_right", "costa_6_left", "costa_6_right", "costa_7_left", "costa_7_right", "costa_8_left", "costa_8_right",
 65    "costa_9_left", "costa_9_right", "costa_10_left", "costa_10_right", "costa_11_left", "costa_11_right",
 66    "costa_12_left", "costa_12_right", "rib_cartilage", "sternum_corpus", "clavicula_left", "clavicula_right",
 67    "scapula_left", "scapula_right", "humerus_left", "humerus_right", "skull", "hip_left", "hip_right", "sacrum",
 68    "femur_left", "femur_right", "heart", "heart_atrium_left", "heart_tissue", "heart_atrium_right",
 69    "heart_myocardium", "heart_ventricle_left", "heart_ventricle_right", "iliac_artery_left", "iliac_artery_right",
 70    "aorta", "iliac_vena_left", "iliac_vena_right", "inferior_vena_cava", "portal_vein_and_splenic_vein",
 71    "celiac_trunk", "lung_lower_lobe_left", "lung_upper_lobe_left", "lung_lower_lobe_right", "lung_middle_lobe_right",
 72    "lung_upper_lobe_right", "bronchus", "trachea", "pulmonary_artery", "cheek_left", "cheek_right", "eyeball_left",
 73    "eyeball_right", "nasal_cavity", "common_carotid_artery_right", "common_carotid_artery_left", "sternum_manubrium",
 74    "internal_carotid_artery_right", "internal_carotid_artery_left", "internal_jugular_vein_right",
 75    "internal_jugular_vein_left", "brachiocephalic_artery", "brachiocephalic_vein_right", "brachiocephalic_vein_left",
 76    "subclavian_artery_right", "subclavian_artery_left",
 77]
 78"""The classes of the DAP Atlas dataset. The label id of a class is its 1-based index."""
 79
 80CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 81"""Mapping from the class name to its label id in the label volumes."""
 82
 83
 84def get_dap_atlas_data(
 85    path: Union[os.PathLike, str], autopet_path: Optional[Union[os.PathLike, str]] = None, download: bool = False
 86) -> Tuple[str, str]:
 87    """Download the DAP Atlas dataset.
 88
 89    Args:
 90        path: Filepath to a folder where the data is downloaded for further processing.
 91        autopet_path: Filepath to the folder where the AutoPET dataset is (or will be) downloaded.
 92            By default, it is downloaded to the 'autopet' sub-folder of `path`.
 93        download: Whether to download the data if it is not present.
 94
 95    Returns:
 96        Filepath where the label volumes are downloaded.
 97        Filepath where the AutoPET data is downloaded.
 98    """
 99    if autopet_path is None:
100        autopet_path = os.path.join(path, "autopet")
101    get_autopet_data(autopet_path, download)
102
103    label_dir = os.path.join(path, "Atlas_final_dataset_V1_533")
104    if not os.path.exists(label_dir):
105        os.makedirs(path, exist_ok=True)
106        zip_path = os.path.join(path, "dap_atlas_masks.zip")
107        util.download_source_gdrive(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
108        util.unzip(zip_path=zip_path, dst=path, remove=False)
109
110    return label_dir, autopet_path
111
112
113def get_dap_atlas_paths(
114    path: Union[os.PathLike, str], autopet_path: Optional[Union[os.PathLike, str]] = None, download: bool = False
115) -> Tuple[List[str], List[str]]:
116    """Get paths to the DAP Atlas data.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        autopet_path: Filepath to the folder where the AutoPET dataset is (or will be) downloaded.
121            By default, it is downloaded to the 'autopet' sub-folder of `path`.
122        download: Whether to download the data if it is not present.
123
124    Returns:
125        List of filepaths for the image data.
126        List of filepaths for the label data.
127    """
128    label_dir, autopet_path = get_dap_atlas_data(path, autopet_path, download)
129
130    # Match each label volume 'AutoPET_<subject id>_<study suffix>.nii.gz' to the CT in the AutoPET study folder
131    # 'PETCT_<subject id>/<study description>-<study suffix>/CT.nii.gz'.
132    autopet_dir = os.path.join(autopet_path, "AutoPET-II", "FDG-PET-CT-Lesions")
133    raw_paths, label_paths, missing = [], [], []
134    for label_path in natsorted(glob(os.path.join(label_dir, "AutoPET_*.nii.gz"))):
135        _, subject_id, study_suffix = os.path.basename(label_path)[:-len(".nii.gz")].split("_")
136        ct_paths = glob(os.path.join(autopet_dir, f"PETCT_{subject_id}", f"*-{study_suffix}", "CT.nii.gz"))
137        if len(ct_paths) > 1:
138            raise RuntimeError(f"Found multiple AutoPET CTs for '{label_path}': {ct_paths}.")
139        if not ct_paths:  # The AutoPET download is incomplete, so this case is skipped.
140            missing.append(os.path.basename(label_path))
141            continue
142        raw_paths.append(ct_paths[0])
143        label_paths.append(label_path)
144
145    if missing:
146        warn(f"Could not find the AutoPET CT for {len(missing)} of {len(missing) + len(raw_paths)} label volumes. "
147             f"These cases are skipped. Is the AutoPET data in '{autopet_path}' complete?")
148
149    if not raw_paths:
150        raise RuntimeError(f"Could not match any DAP Atlas label volume to an AutoPET CT in '{autopet_path}'.")
151
152    assert len(raw_paths) == len(label_paths)
153    return raw_paths, label_paths
154
155
156def get_dap_atlas_dataset(
157    path: Union[os.PathLike, str],
158    patch_shape: Tuple[int, ...],
159    autopet_path: Optional[Union[os.PathLike, str]] = None,
160    resize_inputs: bool = False,
161    download: bool = False,
162    **kwargs
163) -> Dataset:
164    """Get the DAP Atlas dataset for anatomical structure segmentation in whole-body CT.
165
166    Args:
167        path: Filepath to a folder where the data is downloaded for further processing.
168        patch_shape: The patch shape to use for training.
169        autopet_path: Filepath to the folder where the AutoPET dataset is (or will be) downloaded.
170            By default, it is downloaded to the 'autopet' sub-folder of `path`.
171        resize_inputs: Whether to resize inputs to the desired patch shape.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
174
175    Returns:
176        The segmentation dataset.
177    """
178    raw_paths, label_paths = get_dap_atlas_paths(path, autopet_path, download)
179
180    if resize_inputs:
181        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
182        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
183            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
184        )
185
186    return torch_em.default_segmentation_dataset(
187        raw_paths=raw_paths,
188        raw_key="data",
189        label_paths=label_paths,
190        label_key="data",
191        patch_shape=patch_shape,
192        is_seg_dataset=True,
193        **kwargs
194    )
195
196
197def get_dap_atlas_loader(
198    path: Union[os.PathLike, str],
199    batch_size: int,
200    patch_shape: Tuple[int, ...],
201    autopet_path: Optional[Union[os.PathLike, str]] = None,
202    resize_inputs: bool = False,
203    download: bool = False,
204    **kwargs
205) -> DataLoader:
206    """Get the DAP Atlas dataloader for anatomical structure segmentation in whole-body CT.
207
208    Args:
209        path: Filepath to a folder where the data is downloaded for further processing.
210        batch_size: The batch size for training.
211        patch_shape: The patch shape to use for training.
212        autopet_path: Filepath to the folder where the AutoPET dataset is (or will be) downloaded.
213            By default, it is downloaded to the 'autopet' sub-folder of `path`.
214        resize_inputs: Whether to resize inputs to the desired patch shape.
215        download: Whether to download the data if it is not present.
216        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
217
218    Returns:
219        The DataLoader.
220    """
221    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
222    dataset = get_dap_atlas_dataset(path, patch_shape, autopet_path, resize_inputs, download, **ds_kwargs)
223    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)