torch_em.data.datasets.medical.topcow

  1"""The TopCoW dataset contains annotations for the vessel components of the circle of Willis
  2in computed tomography angiography (CTA) and magnetic resonance angiography (MRA).
  3
  4The data was curated for the TopCoW challenge (https://topcow24.grand-challenge.org). This module downloads the
  5public training data release of the second edition, which consists of 250 annotated angiographies (125 CTA and
  6125 MRA of the same 125 patients). The 10 unannotated validation images of the release are not exposed here.
  7The modality is selected with the 'modality' argument ('ct' or 'mr').
  8
  9The multi-class segmentations label 13 vessel components of the circle of Willis, see `LABEL_IDS`. Note that the
 10label id 13 and 14 are not used and that the third A2 segment (an anatomical variant) has the label id 15.
 11
 12The data is located at https://doi.org/10.5281/zenodo.15692630.
 13
 14This dataset is from the publication https://doi.org/10.1056/aidbp2500994.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, List, Optional, Literal
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://zenodo.org/records/15692630/files/TopCoW2024_Data_Release.zip"
 31CHECKSUM = "a23d9d0f05ec439472736f65c47fc7408263eb2f2b01a1d1afc934f2e8f2889b"
 32
 33LABEL_IDS = {
 34    "background": 0,
 35    "BA": 1,
 36    "R-PCA": 2,
 37    "L-PCA": 3,
 38    "R-ICA": 4,
 39    "R-MCA": 5,
 40    "L-ICA": 6,
 41    "L-MCA": 7,
 42    "R-Pcom": 8,
 43    "L-Pcom": 9,
 44    "Acom": 10,
 45    "R-ACA": 11,
 46    "L-ACA": 12,
 47    "3rd-A2": 15,
 48}
 49
 50MODALITIES = ["ct", "mr"]
 51
 52
 53def get_topcow_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 54    """Download the TopCoW dataset.
 55
 56    Args:
 57        path: Filepath to a folder where the data is downloaded for further processing.
 58        download: Whether to download the data if it is not present.
 59
 60    Returns:
 61        Filepath where the data is downloaded.
 62    """
 63    data_dir = os.path.join(path, "TopCoW2024_Data_Release")
 64    if os.path.exists(data_dir):
 65        return data_dir
 66
 67    os.makedirs(path, exist_ok=True)
 68
 69    zip_path = os.path.join(path, "TopCoW2024_Data_Release.zip")
 70    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 71    util.unzip(zip_path=zip_path, dst=path)
 72
 73    return data_dir
 74
 75
 76def get_topcow_paths(
 77    path: Union[os.PathLike, str],
 78    modality: Optional[Literal["ct", "mr"]] = None,
 79    download: bool = False,
 80) -> Tuple[List[str], List[str]]:
 81    """Get paths to the TopCoW data.
 82
 83    Args:
 84        path: Filepath to a folder where the data is downloaded for further processing.
 85        modality: The angiography modality. Either 'ct' or 'mr'. If None, both modalities are returned.
 86        download: Whether to download the data if it is not present.
 87
 88    Returns:
 89        List of filepaths for the image data.
 90        List of filepaths for the label data.
 91    """
 92    data_dir = get_topcow_data(path, download)
 93
 94    if modality is not None and modality not in MODALITIES:
 95        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {MODALITIES}.")
 96
 97    pattern = f"topcow_{'*' if modality is None else modality}_*.nii.gz"
 98    label_paths = natsorted(glob(os.path.join(data_dir, "cow_seg_labelsTr", pattern)))
 99    # The images carry the channel suffix '_0000' of the nnU-Net format, the labels do not.
100    raw_paths = [
101        os.path.join(data_dir, "imagesTr", os.path.basename(p).replace(".nii.gz", "_0000.nii.gz"))
102        for p in label_paths
103    ]
104    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths)
105
106    return raw_paths, label_paths
107
108
109def get_topcow_dataset(
110    path: Union[os.PathLike, str],
111    patch_shape: Tuple[int, ...],
112    modality: Optional[Literal["ct", "mr"]] = None,
113    resize_inputs: bool = False,
114    download: bool = False,
115    **kwargs
116) -> Dataset:
117    """Get the TopCoW dataset for circle of Willis segmentation.
118
119    Args:
120        path: Filepath to a folder where the data is downloaded for further processing.
121        patch_shape: The patch shape to use for training.
122        modality: The angiography modality. Either 'ct' or 'mr'. If None, both modalities are returned.
123        resize_inputs: Whether to resize inputs to the desired patch shape.
124        download: Whether to download the data if it is not present.
125        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
126
127    Returns:
128        The segmentation dataset.
129    """
130    raw_paths, label_paths = get_topcow_paths(path, modality, download)
131
132    if resize_inputs:
133        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
134        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
135            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
136        )
137
138    return torch_em.default_segmentation_dataset(
139        raw_paths=raw_paths,
140        raw_key="data",
141        label_paths=label_paths,
142        label_key="data",
143        patch_shape=patch_shape,
144        is_seg_dataset=True,
145        **kwargs
146    )
147
148
149def get_topcow_loader(
150    path: Union[os.PathLike, str],
151    batch_size: int,
152    patch_shape: Tuple[int, ...],
153    modality: Optional[Literal["ct", "mr"]] = None,
154    resize_inputs: bool = False,
155    download: bool = False,
156    **kwargs
157) -> DataLoader:
158    """Get the TopCoW dataloader for circle of Willis segmentation.
159
160    Args:
161        path: Filepath to a folder where the data is downloaded for further processing.
162        batch_size: The batch size for training.
163        patch_shape: The patch shape to use for training.
164        modality: The angiography modality. Either 'ct' or 'mr'. If None, both modalities are returned.
165        resize_inputs: Whether to resize inputs to the desired patch shape.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
168
169    Returns:
170        The DataLoader.
171    """
172    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
173    dataset = get_topcow_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
174    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)