torch_em.data.datasets.medical.aortaseg24

  1"""The AortaSeg dataset contains annotations for multi-class segmentation of the aortic branches and zones
  2in computed tomography angiography (CTA) scans.
  3
  4It comprises the training set of the AortaSeg24 challenge (https://aortaseg24.grand-challenge.org):
  550 CTA volumes of patients with uncomplicated type B aortic dissection, resampled to an isotropic resolution
  6of 1mm, with 23 annotated aortic zones and branches.
  7
  8NOTE: The label legend is as follows (see `AORTIC_SEGMENTS` and `LABEL_IDS`; 1 = Zone 0, 2 = Innominate Artery,
  9..., 23 = Zone 11 L). The ids were taken from the official evaluation code of the challenge
 10(https://github.com/ImranNust/AortaSeg24/blob/main/evaluation_docker_for_validation_phase/evaluate.py),
 11which one-hot encodes the labels with 24 classes and reports the per-class dice in this order.
 12
 13NOTE: The dataset requires registration and cannot be downloaded automatically. Please follow these steps:
 14- Visit https://aortaseg24.grand-challenge.org/dataset-access-information/ and complete the dataset access
 15  agreement form via the DocuSign link given there. The approval may take up to 24 hours.
 16- Join the challenge at https://aortaseg24.grand-challenge.org/ and request access to the dataset on the
 17  dataset page. You will then receive the link to the Dropbox folder with the data.
 18- Download the training images and masks and place them at '<path>', such that
 19  '<path>/images/subject001_CTA.mha' and '<path>/masks/subject001_label.mha' exist.
 20
 21The dataset is located at https://aortaseg24.grand-challenge.org/dataset-access-information/.
 22
 23This dataset is from the publication https://doi.org/10.1016/j.media.2026.104188.
 24Please cite it if you use this dataset in your research.
 25
 26NOTE: Reading the MetaImage (.mha) volumes requires 'SimpleITK'. Install it with 'pip install SimpleITK'.
 27"""
 28
 29import os
 30from glob import glob
 31from natsort import natsorted
 32from typing import Union, Tuple, List
 33
 34from torch.utils.data import Dataset, DataLoader
 35
 36import torch_em
 37
 38from .. import util
 39
 40
 41AORTIC_SEGMENTS = [
 42    "Zone_0", "Innominate_Artery", "Zone_1", "Left_Common_Carotid", "Zone_2", "Left_Subclavian_Artery", "Zone_3",
 43    "Zone_4", "Zone_5", "Zone_6", "Celiac_Artery", "Zone_7", "SMA", "Zone_8", "Right_Renal_Artery",
 44    "Left_Renal_Artery", "Zone_9", "Zone_10_R", "Zone_10_L", "Right_Internal_Iliac_Artery",
 45    "Left_Internal_Iliac_Artery", "Zone_11_R", "Zone_11_L",
 46]
 47
 48LABEL_IDS = {"background": 0, **{name: i + 1 for i, name in enumerate(AORTIC_SEGMENTS)}}
 49
 50
 51def get_aortaseg24_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 52    """Obtain the AortaSeg dataset.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        download: Whether to download the data if it is not present.
 57
 58    Returns:
 59        Filepath where the data is stored.
 60    """
 61    if download:
 62        msg = "Download is set to True, but 'torch_em' cannot download this dataset. "
 63        msg += "See 'torch_em.data.datasets.medical.aortaseg24' for the manual download instructions."
 64        raise NotImplementedError(msg)
 65
 66    # The data is either placed directly in 'path' or in a subfolder, e.g. named after the downloaded archive.
 67    image_dirs = [p for p in glob(os.path.join(path, "**", "images"), recursive=True) if os.path.isdir(p)]
 68    if len(image_dirs) == 0:
 69        raise FileNotFoundError(
 70            f"It's expected to place the downloaded AortaSeg24 training data at '{path}'. "
 71            "See 'torch_em.data.datasets.medical.aortaseg24' for the manual download instructions."
 72        )
 73
 74    return os.path.split(image_dirs[0])[0]
 75
 76
 77def get_aortaseg24_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 78    """Get paths to the AortaSeg data.
 79
 80    Args:
 81        path: Filepath to a folder where the data is downloaded for further processing.
 82        download: Whether to download the data if it is not present.
 83
 84    Returns:
 85        List of filepaths for the image data.
 86        List of filepaths for the label data.
 87    """
 88    data_dir = get_aortaseg24_data(path, download)
 89
 90    raw_paths = natsorted(glob(os.path.join(data_dir, "images", "*_CTA.mha")))
 91    label_paths = [
 92        os.path.join(data_dir, "masks", f"{os.path.basename(p)[:-len('_CTA.mha')]}_label.mha") for p in raw_paths
 93    ]
 94    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 95
 96    return raw_paths, label_paths
 97
 98
 99def get_aortaseg24_dataset(
100    path: Union[os.PathLike, str],
101    patch_shape: Tuple[int, ...],
102    resize_inputs: bool = False,
103    download: bool = False,
104    **kwargs
105) -> Dataset:
106    """Get the AortaSeg dataset for aortic branch and zone segmentation.
107
108    Args:
109        path: Filepath to a folder where the data is downloaded for further processing.
110        patch_shape: The patch shape to use for training.
111        resize_inputs: Whether to resize inputs to the desired patch shape.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
114
115    Returns:
116        The segmentation dataset.
117    """
118    raw_paths, label_paths = get_aortaseg24_paths(path, download)
119
120    if resize_inputs:
121        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
122        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
123            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
124        )
125
126    return torch_em.default_segmentation_dataset(
127        raw_paths=raw_paths,
128        raw_key=None,
129        label_paths=label_paths,
130        label_key=None,
131        patch_shape=patch_shape,
132        is_seg_dataset=True,
133        **kwargs
134    )
135
136
137def get_aortaseg24_loader(
138    path: Union[os.PathLike, str],
139    batch_size: int,
140    patch_shape: Tuple[int, ...],
141    resize_inputs: bool = False,
142    download: bool = False,
143    **kwargs
144) -> DataLoader:
145    """Get the AortaSeg dataloader for aortic branch and zone segmentation.
146
147    Args:
148        path: Filepath to a folder where the data is downloaded for further processing.
149        batch_size: The batch size for training.
150        patch_shape: The patch shape to use for training.
151        resize_inputs: Whether to resize inputs to the desired patch shape.
152        download: Whether to download the data if it is not present.
153        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
154
155    Returns:
156        The DataLoader.
157    """
158    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
159    dataset = get_aortaseg24_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
160    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)