torch_em.data.datasets.medical.topaneu

  1"""The TopAneu dataset contains annotations for vessel-specific intracranial aneurysm
  2classification and segmentation in computed tomography angiography (CTA) and magnetic
  3resonance angiography (MRA).
  4
  5The data was curated for the TopAneu 2026 challenge (https://topaneu-26.grand-challenge.org), which extends
  6the TopBrain / TopCoW vessel anatomy challenges (`torch_em.data.datasets.medical.topbrain`,
  7`torch_em.data.datasets.medical.topcow`) to vessel-specific aneurysm classification and segmentation. This
  8module downloads the training data release (batch-1 and batch-2, published June 15 and July 31 2026), which
  9consists of 415 annotated angiographies (415 scans from 408 patients) from Lausanne University Hospital
 10(CHUV), Geneva University Hospitals (HUG), Mie Chuo Medical Center, and public data reused from INSTED and
 11the Lausanne TOF-MRA aneurysm cohort on OpenNeuro. The modality is selected with the 'modality' argument
 12('ct' or 'mr').
 13
 14Three types of voxel-level annotations are provided, selected with the 'label_choice' argument:
 15- 'location': multiclass aneurysm segmentation, where each aneurysm voxel is assigned to one of 52
 16  vessel-specific location classes (laterality x anatomical position), see the 'location_mapping.json'
 17  file of the downloaded data.
 18- 'type': multiclass aneurysm segmentation by morphological type (saccular, dissecting, fusiform), see the
 19  'type_mapping.json' file of the downloaded data.
 20- 'vessel': the vessel anatomy mask predicted by the TopBrain organizer model, see the 'vessel_mapping.json'
 21  file of the downloaded data. This is a silver-standard annotation, not a manual ground-truth.
 22
 23The data is hosted in a public SWITCHdrive share at https://drive.switch.ch/index.php/s/O36U43RkChkNcHd (see
 24also https://topaneu-26.grand-challenge.org/data). It is released under a non-commercial open-use license
 25that requires attribution, see the 'Terms_of_use.txt' file of the downloaded data for the full terms.
 26
 27The challenge design is described in https://doi.org/10.5281/zenodo.19848807.
 28Please cite the TopAneu 2026 challenge if you use this dataset in your research.
 29"""
 30
 31import os
 32import re
 33from glob import glob
 34from tqdm import tqdm
 35from natsort import natsorted
 36from urllib.parse import quote, unquote
 37from typing import Union, Tuple, List, Optional, Literal
 38
 39import requests
 40
 41from torch.utils.data import Dataset, DataLoader
 42
 43import torch_em
 44
 45from .. import util
 46
 47
 48SHARE_TOKEN = "O36U43RkChkNcHd"
 49
 50URL = f"https://drive.switch.ch/index.php/s/{SHARE_TOKEN}"
 51
 52WEBDAV_URL = "https://drive.switch.ch/public.php/webdav"
 53
 54# The files are downloaded individually from the public share, so there is no checksum for a single archive.
 55CHECKSUM = None
 56
 57MODALITIES = ["ct", "mr"]
 58
 59LABEL_CHOICES = {"location": "location_masks", "type": "type_masks", "vessel": "vessel_masks"}
 60
 61
 62def _list_share_folder(folder):
 63    """List the file names in a folder of the public share via the WebDAV endpoint of Nextcloud.
 64
 65    Public shares are accessed by using the share token as the user name and an empty password.
 66    """
 67    response = requests.request(
 68        "PROPFIND", f"{WEBDAV_URL}/{quote(folder)}/", headers={"Depth": "1"}, auth=(SHARE_TOKEN, "")
 69    )
 70    response.raise_for_status()
 71
 72    fnames = []
 73    for href in re.findall(r"<d:href>(.*?)</d:href>", response.text):
 74        fname = unquote(href).rstrip("/").split("/")[-1]
 75        if fname.endswith(".nii.gz"):
 76            fnames.append(fname)
 77    return natsorted(fnames)
 78
 79
 80def _download_share_folder(folder, dst, download):
 81    """Download all nifti files of a folder of the public share into `dst`."""
 82    if os.path.exists(dst):
 83        return
 84
 85    if not download:
 86        raise RuntimeError(f"Cannot find the data at {dst}, but download was set to False.")
 87
 88    tmp_dir = f"{dst}.tmp"
 89    os.makedirs(tmp_dir, exist_ok=True)
 90    fnames = _list_share_folder(folder)
 91    for fname in tqdm(fnames, desc=f"Download {len(fnames)} files from '{folder}'"):
 92        out_path = os.path.join(tmp_dir, fname)
 93        if os.path.exists(out_path):
 94            continue
 95        url = f"{URL}/download?path={quote('/' + folder)}&files={quote(fname)}"
 96        util.download_source(path=out_path, url=url, download=download, checksum=CHECKSUM)
 97
 98    os.rename(tmp_dir, dst)
 99
100
101def get_topaneu_data(
102    path: Union[os.PathLike, str], label_choice: Literal["location", "type", "vessel"] = "location",
103    download: bool = False,
104) -> str:
105    """Download the TopAneu dataset.
106
107    Args:
108        path: Filepath to a folder where the data is downloaded for further processing.
109        label_choice: The choice of segmentation target. One of 'location', 'type' or 'vessel'.
110        download: Whether to download the data if it is not present.
111
112    Returns:
113        Filepath where the data is downloaded.
114    """
115    if label_choice not in LABEL_CHOICES:
116        raise ValueError(f"'{label_choice}' is not a valid label choice. Please choose one of {list(LABEL_CHOICES)}.")
117
118    data_dir = os.path.join(path, "topaneu")
119    os.makedirs(path, exist_ok=True)
120
121    _download_share_folder("images", os.path.join(data_dir, "images"), download)
122    _download_share_folder(LABEL_CHOICES[label_choice], os.path.join(data_dir, LABEL_CHOICES[label_choice]), download)
123
124    return data_dir
125
126
127def get_topaneu_paths(
128    path: Union[os.PathLike, str],
129    label_choice: Literal["location", "type", "vessel"] = "location",
130    modality: Optional[Literal["ct", "mr"]] = None,
131    download: bool = False,
132) -> Tuple[List[str], List[str]]:
133    """Get paths to the TopAneu data.
134
135    Args:
136        path: Filepath to a folder where the data is downloaded for further processing.
137        label_choice: The choice of segmentation target. One of 'location', 'type' or 'vessel'.
138        modality: The angiography modality. Either 'ct' or 'mr'. If None, both modalities are returned.
139        download: Whether to download the data if it is not present.
140
141    Returns:
142        List of filepaths for the image data.
143        List of filepaths for the label data.
144    """
145    data_dir = get_topaneu_data(path, label_choice, download)
146
147    if modality is not None and modality not in MODALITIES:
148        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {MODALITIES}.")
149
150    pattern = "topaneu_*.nii.gz" if modality is None else f"topaneu_*_{modality}_*.nii.gz"
151    label_paths = natsorted(glob(os.path.join(data_dir, LABEL_CHOICES[label_choice], pattern)))
152    # The images carry the channel suffix '_0000' of the nnU-Net format, the labels do not.
153    raw_paths = [
154        os.path.join(data_dir, "images", os.path.basename(p).replace(".nii.gz", "_0000.nii.gz"))
155        for p in label_paths
156    ]
157    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths)
158
159    return raw_paths, label_paths
160
161
162def get_topaneu_dataset(
163    path: Union[os.PathLike, str],
164    patch_shape: Tuple[int, ...],
165    label_choice: Literal["location", "type", "vessel"] = "location",
166    modality: Optional[Literal["ct", "mr"]] = None,
167    resize_inputs: bool = False,
168    download: bool = False,
169    **kwargs
170) -> Dataset:
171    """Get the TopAneu dataset for vessel-specific intracranial aneurysm segmentation.
172
173    Args:
174        path: Filepath to a folder where the data is downloaded for further processing.
175        patch_shape: The patch shape to use for training.
176        label_choice: The choice of segmentation target. One of 'location', 'type' or 'vessel'.
177        modality: The angiography modality. Either 'ct' or 'mr'. If None, both modalities are returned.
178        resize_inputs: Whether to resize inputs to the desired patch shape.
179        download: Whether to download the data if it is not present.
180        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
181
182    Returns:
183        The segmentation dataset.
184    """
185    raw_paths, label_paths = get_topaneu_paths(path, label_choice, modality, download)
186
187    if resize_inputs:
188        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
189        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
190            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
191        )
192
193    return torch_em.default_segmentation_dataset(
194        raw_paths=raw_paths,
195        raw_key="data",
196        label_paths=label_paths,
197        label_key="data",
198        patch_shape=patch_shape,
199        is_seg_dataset=True,
200        **kwargs
201    )
202
203
204def get_topaneu_loader(
205    path: Union[os.PathLike, str],
206    batch_size: int,
207    patch_shape: Tuple[int, ...],
208    label_choice: Literal["location", "type", "vessel"] = "location",
209    modality: Optional[Literal["ct", "mr"]] = None,
210    resize_inputs: bool = False,
211    download: bool = False,
212    **kwargs
213) -> DataLoader:
214    """Get the TopAneu dataloader for vessel-specific intracranial aneurysm segmentation.
215
216    Args:
217        path: Filepath to a folder where the data is downloaded for further processing.
218        batch_size: The batch size for training.
219        patch_shape: The patch shape to use for training.
220        label_choice: The choice of segmentation target. One of 'location', 'type' or 'vessel'.
221        modality: The angiography modality. Either 'ct' or 'mr'. If None, both modalities are returned.
222        resize_inputs: Whether to resize inputs to the desired patch shape.
223        download: Whether to download the data if it is not present.
224        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
225
226    Returns:
227        The DataLoader.
228    """
229    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
230    dataset = get_topaneu_dataset(path, patch_shape, label_choice, modality, resize_inputs, download, **ds_kwargs)
231    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)