torch_em.data.datasets.medical.isles2024

  1"""The ISLES 2024 dataset contains annotations for ischemic stroke lesion and large vessel occlusion
  2segmentation in a multimodal, longitudinal collection of brain CT and MRI scans.
  3
  4The dataset consists of 149 acute ischemic stroke cases, organized according to the BIDS standard. For each
  5case, the following data is provided: admission non-contrast CT (NCCT), CT angiography (CTA) and 4D CT
  6perfusion (CTP) with its derived perfusion maps (Tmax, CBF, CBV, MTT), as well as follow-up MRI (DWI, ADC).
  7This module uses the 'derivatives' folder, in which all modalities are linearly co-registered to the NCCT
  8space, so that raw and label volumes of the same case share a common voxel grid.
  9
 10Two segmentation targets are provided, selected with the 'label_choice' argument:
 11- 'lesion': the binary infarct mask, derived from the follow-up MRI (session 'ses-02'). The corresponding
 12  modalities are 'dwi' and 'adc'.
 13- 'lvo': the binary large vessel occlusion mask, derived from the admission CTA (session 'ses-01'). The
 14  corresponding modalities are 'ncct', 'cta', 'tmax', 'mtt', 'cbf' and 'cbv'.
 15The modality is selected with the 'modality' argument. If it is not specified, all modalities of the chosen
 16'label_choice' are stacked as channels of the raw input.
 17
 18The data is located at https://doi.org/10.5281/zenodo.16813698.
 19
 20NOTE: The archive is distributed as a single ~92 GB 7z file, which requires the 'p7zip' CLI to extract
 21(see `torch_em.data.datasets.util.unzip_7z`).
 22
 23This dataset is from the publication https://doi.org/10.48550/arXiv.2408.11142.
 24Please cite it if you use this dataset in your research.
 25"""
 26
 27import os
 28from glob import glob
 29from natsort import natsorted
 30from typing import Union, Tuple, List, Optional, Literal
 31
 32from torch.utils.data import Dataset, DataLoader
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39URL = "https://zenodo.org/records/16813698/files/train.7z"
 40CHECKSUM = "038920e4dc2011a3f47b8bb8421c67e36d07f1d84f1ba442563077480f75d129"
 41
 42LABEL_CHOICES = ["lesion", "lvo"]
 43
 44MODALITIES = {
 45    "lesion": ["dwi", "adc"],
 46    "lvo": ["ncct", "cta", "tmax", "mtt", "cbf", "cbv"],
 47}
 48
 49SESSIONS = {"lesion": "ses-02", "lvo": "ses-01"}
 50
 51PERFUSION_MAPS = ["tmax", "mtt", "cbf", "cbv"]
 52
 53
 54def get_isles2024_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 55    """Download the ISLES 2024 dataset.
 56
 57    Args:
 58        path: Filepath to a folder where the data is downloaded for further processing.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        Filepath where the data is downloaded.
 63    """
 64    data_dir = os.path.join(path, "train")
 65    if os.path.exists(os.path.join(data_dir, "derivatives")):
 66        return data_dir
 67
 68    os.makedirs(path, exist_ok=True)
 69
 70    archive_path = os.path.join(path, "train.7z")
 71    util.download_source(path=archive_path, url=URL, download=download, checksum=CHECKSUM)
 72    util.unzip_7z(path_7z=archive_path, dst=path, remove=False)
 73
 74    assert os.path.exists(os.path.join(data_dir, "derivatives")), \
 75        f"The extraction of the ISLES 2024 archive did not create '{data_dir}'."
 76    return data_dir
 77
 78
 79def _get_raw_path(label_path: str, label_choice: str, modality: str) -> str:
 80    case_dir = os.path.dirname(label_path)
 81    base_name = os.path.basename(label_path).replace(f"_{label_choice}-msk.nii.gz", f"_{modality}.nii.gz")
 82    if modality in PERFUSION_MAPS:
 83        return os.path.join(case_dir, "perfusion-maps", base_name)
 84    return os.path.join(case_dir, base_name)
 85
 86
 87def get_isles2024_paths(
 88    path: Union[os.PathLike, str],
 89    label_choice: Literal["lesion", "lvo"] = "lesion",
 90    modality: Optional[str] = None,
 91    download: bool = False,
 92) -> Tuple[List[str], List[str]]:
 93    """Get paths to the ISLES 2024 data.
 94
 95    Args:
 96        path: Filepath to a folder where the data is downloaded for further processing.
 97        label_choice: The choice of segmentation target. Either 'lesion' (infarct) or 'lvo'
 98            (large vessel occlusion).
 99        modality: The choice of imaging modality. See `MODALITIES` for the valid choices per 'label_choice'.
100            If None, all modalities of the chosen 'label_choice' are returned.
101        download: Whether to download the data if it is not present.
102
103    Returns:
104        List of filepaths for the image data.
105        List of filepaths for the label data.
106    """
107    if label_choice not in LABEL_CHOICES:
108        raise ValueError(f"'{label_choice}' is not a valid label choice. Please choose one of {LABEL_CHOICES}.")
109
110    valid_modalities = MODALITIES[label_choice]
111    if modality is not None and modality not in valid_modalities:
112        raise ValueError(f"'{modality}' is not a valid modality for '{label_choice}'. Choose one of {valid_modalities}.")  # noqa
113
114    data_dir = get_isles2024_data(path, download)
115    session = SESSIONS[label_choice]
116
117    label_paths = natsorted(
118        glob(os.path.join(data_dir, "derivatives", "sub-*", session, f"*_{label_choice}-msk.nii.gz"))
119    )
120    assert len(label_paths) > 0, f"Could not find any '{label_choice}' labels in '{data_dir}'."
121
122    modalities = valid_modalities if modality is None else [modality]
123    if modality is None:
124        image_paths = [tuple(_get_raw_path(lp, label_choice, m) for m in modalities) for lp in label_paths]
125        for paths_per_case in image_paths:
126            assert all(os.path.exists(p) for p in paths_per_case)
127    else:
128        image_paths = [_get_raw_path(lp, label_choice, modality) for lp in label_paths]
129        assert all(os.path.exists(p) for p in image_paths)
130
131    return image_paths, label_paths
132
133
134def get_isles2024_dataset(
135    path: Union[os.PathLike, str],
136    patch_shape: Tuple[int, ...],
137    label_choice: Literal["lesion", "lvo"] = "lesion",
138    modality: Optional[str] = None,
139    resize_inputs: bool = False,
140    download: bool = False,
141    **kwargs
142) -> Dataset:
143    """Get the ISLES 2024 dataset for ischemic stroke lesion and large vessel occlusion segmentation.
144
145    Args:
146        path: Filepath to a folder where the data is downloaded for further processing.
147        patch_shape: The patch shape to use for training.
148        label_choice: The choice of segmentation target. Either 'lesion' (infarct) or 'lvo'
149            (large vessel occlusion).
150        modality: The choice of imaging modality. See `MODALITIES` for the valid choices per 'label_choice'.
151            If None, all modalities of the chosen 'label_choice' are returned.
152        resize_inputs: Whether to resize inputs to the desired patch shape.
153        download: Whether to download the data if it is not present.
154        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
155
156    Returns:
157        The segmentation dataset.
158    """
159    image_paths, label_paths = get_isles2024_paths(path, label_choice, modality, download)
160
161    if resize_inputs:
162        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
163        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
164            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
165        )
166
167    dataset = torch_em.default_segmentation_dataset(
168        raw_paths=image_paths,
169        raw_key="data",
170        label_paths=label_paths,
171        label_key="data",
172        patch_shape=patch_shape,
173        with_channels=modality is None,
174        is_seg_dataset=True,
175        **kwargs
176    )
177    if "sampler" in kwargs:
178        for ds in dataset.datasets:
179            ds.max_sampling_attempts = 5000
180
181    return dataset
182
183
184def get_isles2024_loader(
185    path: Union[os.PathLike, str],
186    batch_size: int,
187    patch_shape: Tuple[int, ...],
188    label_choice: Literal["lesion", "lvo"] = "lesion",
189    modality: Optional[str] = None,
190    resize_inputs: bool = False,
191    download: bool = False,
192    **kwargs
193) -> DataLoader:
194    """Get the ISLES 2024 dataloader for ischemic stroke lesion and large vessel occlusion segmentation.
195
196    Args:
197        path: Filepath to a folder where the data is downloaded for further processing.
198        batch_size: The batch size for training.
199        patch_shape: The patch shape to use for training.
200        label_choice: The choice of segmentation target. Either 'lesion' (infarct) or 'lvo'
201            (large vessel occlusion).
202        modality: The choice of imaging modality. See `MODALITIES` for the valid choices per 'label_choice'.
203            If None, all modalities of the chosen 'label_choice' are returned.
204        resize_inputs: Whether to resize inputs to the desired patch shape.
205        download: Whether to download the data if it is not present.
206        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
207
208    Returns:
209        The DataLoader.
210    """
211    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
212    dataset = get_isles2024_dataset(path, patch_shape, label_choice, modality, resize_inputs, download, **ds_kwargs)
213    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)