torch_em.data.datasets.medical.brats24

  1"""The BraTS 2024 dataset contains annotations for the sub-regions of post-treatment adult diffuse
  2glioma in multi-modal brain MRI.
  3
  4It is the adult glioma segmentation task of the 2024 Brain Tumor Segmentation (BraTS) challenge
  5(https://www.synapse.org/Synapse:syn53708126). Unlike the BraTS 2023 release (see `brats.py`), this is an
  6entirely new dataset of exclusively post-treatment studies, with 1621 training studies. Each study provides
  7four co-registered, skull-stripped and interpolated sequences, which can be selected with the 'modality'
  8argument: a native T1-weighted scan ('t1n'), a post-contrast T1-weighted scan ('t1c'), a T2-weighted scan
  9('t2w') and a T2 FLAIR scan ('t2f').
 10
 11The label ids are described in `LABEL_IDS`: 0 = background, 1 = non-enhancing tumor core (NETC), 2 =
 12surrounding non-enhancing FLAIR hyperintensity (SNFH), 3 = enhancing tissue (ET), 4 = resection cavity (RC).
 13NOTE: These ids differ from the BraTS 2023 (and earlier) releases, which do not have a resection cavity
 14class and instead use the id 1 for the necrotic tumor core.
 15
 16Evaluation is not done on the sub-regions themselves, but on the nested regions that they form, which can be
 17selected with the 'region' argument (see `REGIONS`): the whole tumor (the union of the non-enhancing tumor
 18core, the FLAIR hyperintensity and the enhancing tissue, excluding the resection cavity), the tumor core
 19(the union of the non-enhancing tumor core and the enhancing tissue) and the enhancing tissue. The
 20individual sub-regions, including the resection cavity, can also be selected as a binary target with this
 21argument. By default, the sub-region ids are returned as they are. This follows the official evaluation
 22script at https://github.com/rachitsaluja/BraTS-2024-Metrics.
 23
 24NOTE: The official data at https://www.synapse.org/Synapse:syn53708126 is only handed out to registered
 25participants, so this module downloads a public mirror of the BraTS 2024 adult glioma training set at
 26https://huggingface.co/datasets/Spirit-26/BraTS-2024-Complete. If the official release is extracted into
 27the folder passed as 'path', so that files such as
 28'<path>/**/BraTS-GLI-00005-100/BraTS-GLI-00005-100-t2f.nii.gz' exist, it is used instead of the mirror.
 29
 30The scans are used as nifti volumes directly (the key is 'data'). They are loaded with the axis order
 31reversed with respect to the nifti file, i.e. (Z, Y, X), so that a 2d patch shape selects axial slices.
 32
 33This dataset is from the publication https://doi.org/10.48550/arXiv.2405.18368.
 34Please cite it if you use this dataset in your research.
 35"""
 36
 37import os
 38import json
 39from glob import glob
 40from tqdm import tqdm
 41from natsort import natsorted
 42from typing import Union, Tuple, List, Optional, Literal
 43
 44import numpy as np
 45
 46from torch.utils.data import Dataset, DataLoader
 47
 48import torch_em
 49
 50from .. import util
 51
 52
 53FOLDER_NAME = "BraTS-GLI"
 54
 55URL_BASE = f"https://huggingface.co/datasets/Spirit-26/BraTS-2024-Complete/resolve/main/{FOLDER_NAME}/train"
 56
 57API_URL = f"https://huggingface.co/api/datasets/Spirit-26/BraTS-2024-Complete/tree/main/{FOLDER_NAME}/train"
 58
 59LABEL_IDS = {
 60    "background": 0,
 61    "non_enhancing_tumor_core": 1,
 62    "surrounding_flair_hyperintensity": 2,
 63    "enhancing_tissue": 3,
 64    "resection_cavity": 4,
 65}
 66
 67# The nested tumor regions that the challenge evaluates, and the sub-regions they are made of.
 68REGIONS = {
 69    "whole_tumor": (1, 2, 3),
 70    "tumor_core": (1, 3),
 71    "enhancing_tissue": (3,),
 72    "surrounding_flair_hyperintensity": (2,),
 73    "non_enhancing_tumor_core": (1,),
 74    "resection_cavity": (4,),
 75}
 76
 77MODALITIES = ["t1n", "t1c", "t2w", "t2f"]
 78
 79N_SUBJECTS = 1621
 80
 81N_RETRIES = 5
 82
 83
 84class RegionTransform:
 85    """Transform the BraTS 2024 sub-region ids into a binary mask for one of the tumor regions.
 86
 87    Args:
 88        region: The name of the tumor region, see `REGIONS`.
 89    """
 90    def __init__(self, region: str):
 91        self.region = region
 92
 93    def __call__(self, labels: np.ndarray) -> np.ndarray:
 94        """Apply the transform.
 95
 96        Args:
 97            labels: The sub-region ids.
 98
 99        Returns:
100            The binary mask of the tumor region.
101        """
102        return np.isin(labels, REGIONS[self.region]).astype("uint8")
103
104
105def _get_subject_ids(path, download):
106    """List the studies of the mirror via the huggingface API and cache the listing next to the data."""
107    listing_path = os.path.join(path, "subject_ids.json")
108    if os.path.exists(listing_path):
109        with open(listing_path, "r") as f:
110            return json.load(f)
111
112    if not download:
113        raise RuntimeError(f"Cannot find the data at '{path}', but download was set to False.")
114
115    import requests
116
117    subject_ids, cursor = [], None
118    while True:
119        params = {"limit": 1000}
120        if cursor is not None:
121            params["cursor"] = cursor
122
123        response = requests.get(API_URL, params=params)
124        response.raise_for_status()
125        subject_ids.extend(os.path.basename(entry["path"]) for entry in response.json())
126
127        link = response.headers.get("Link", "")
128        if 'rel="next"' not in link:
129            break
130        cursor = link.split("cursor=")[1].split("&")[0].split(">")[0]
131
132    subject_ids = natsorted(subject_ids)
133    assert len(subject_ids) == N_SUBJECTS, f"Expected {N_SUBJECTS} studies in the mirror, got {len(subject_ids)}."
134
135    with open(listing_path, "w") as f:
136        json.dump(subject_ids, f)
137
138    return subject_ids
139
140
141def _find_data(path, modality):
142    """Find the studies on disk, both for the official release and for the mirror downloaded by this module."""
143    pattern = os.path.join(path, "**", "BraTS-GLI-*", f"BraTS-GLI-*-{modality}.nii.gz")
144    raw_paths = natsorted(glob(pattern, recursive=True))
145    label_paths = [p.replace(f"-{modality}.nii.gz", "-seg.nii.gz") for p in raw_paths]
146
147    keep = [i for i, p in enumerate(label_paths) if os.path.exists(p)]
148    return [raw_paths[i] for i in keep], [label_paths[i] for i in keep]
149
150
151def _download_volumes(path, modality, download):
152    raw_paths, label_paths = [], []
153    for subject_id in tqdm(_get_subject_ids(path, download), desc="Downloading the BraTS 2024 studies"):
154        subject_dir = os.path.join(path, FOLDER_NAME, subject_id)
155        os.makedirs(subject_dir, exist_ok=True)
156
157        for suffix in [modality, "seg"]:
158            fname = f"{subject_id}-{suffix}.nii.gz"
159            fpath = os.path.join(subject_dir, fname)
160            # The mirror is fetched file by file, so a transient error is retried instead of failing the download.
161            for attempt in range(N_RETRIES):
162                try:
163                    util.download_source(path=fpath, url=f"{URL_BASE}/{subject_id}/{fname}", download=download)
164                    break
165                except Exception:
166                    if attempt == N_RETRIES - 1:
167                        raise
168
169            (label_paths if suffix == "seg" else raw_paths).append(fpath)
170
171    return raw_paths, label_paths
172
173
174def get_brats24_data(
175    path: Union[os.PathLike, str],
176    modality: Literal["t1n", "t1c", "t2w", "t2f"] = "t2f",
177    download: bool = False,
178) -> Tuple[List[str], List[str]]:
179    """Download the BraTS 2024 post-treatment adult glioma dataset.
180
181    Only the requested modality and the annotations are downloaded, since the studies are fetched study
182    by study from the mirror.
183
184    Args:
185        path: Filepath to a folder where the data is downloaded for further processing.
186        modality: The MRI sequence. Either 't1n', 't1c', 't2w' or 't2f'.
187        download: Whether to download the data if it is not present.
188
189    Returns:
190        List of filepaths for the image data.
191        List of filepaths for the label data.
192    """
193    if modality not in MODALITIES:
194        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {MODALITIES}.")
195
196    os.makedirs(path, exist_ok=True)
197
198    raw_paths, label_paths = _find_data(path, modality)
199    if len(raw_paths) == N_SUBJECTS:
200        return raw_paths, label_paths
201
202    return _download_volumes(path, modality, download)
203
204
205def get_brats24_paths(
206    path: Union[os.PathLike, str],
207    modality: Literal["t1n", "t1c", "t2w", "t2f"] = "t2f",
208    download: bool = False,
209) -> Tuple[List[str], List[str]]:
210    """Get paths to the BraTS 2024 post-treatment adult glioma data.
211
212    Args:
213        path: Filepath to a folder where the data is downloaded for further processing.
214        modality: The MRI sequence. Either 't1n', 't1c', 't2w' or 't2f'.
215        download: Whether to download the data if it is not present.
216
217    Returns:
218        List of filepaths for the image data.
219        List of filepaths for the label data.
220    """
221    raw_paths, label_paths = get_brats24_data(path, modality, download)
222    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0, f"Could not find the studies in '{path}'."
223    return raw_paths, label_paths
224
225
226def get_brats24_dataset(
227    path: Union[os.PathLike, str],
228    patch_shape: Tuple[int, ...],
229    modality: Literal["t1n", "t1c", "t2w", "t2f"] = "t2f",
230    region: Optional[Literal[
231        "whole_tumor", "tumor_core", "enhancing_tissue",
232        "surrounding_flair_hyperintensity", "non_enhancing_tumor_core", "resection_cavity",
233    ]] = None,
234    resize_inputs: bool = False,
235    download: bool = False,
236    **kwargs
237) -> Dataset:
238    """Get the BraTS 2024 post-treatment adult glioma dataset for brain tumor segmentation.
239
240    Args:
241        path: Filepath to a folder where the data is downloaded for further processing.
242        patch_shape: The patch shape to use for training.
243        modality: The MRI sequence. Either 't1n', 't1c', 't2w' or 't2f'.
244        region: The tumor region to use as a binary target, see `REGIONS`. If None, the sub-region ids are used.
245        resize_inputs: Whether to resize inputs to the desired patch shape.
246        download: Whether to download the data if it is not present.
247        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
248
249    Returns:
250        The segmentation dataset.
251    """
252    if region is not None and region not in REGIONS:
253        raise ValueError(f"'{region}' is not a valid region. Please choose one of {list(REGIONS.keys())}.")
254
255    raw_paths, label_paths = get_brats24_paths(path, modality, download)
256
257    if region is not None:
258        kwargs = util.update_kwargs(kwargs, "label_transform", RegionTransform(region))
259
260    if resize_inputs:
261        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
262        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
263            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
264        )
265
266    return torch_em.default_segmentation_dataset(
267        raw_paths=raw_paths,
268        raw_key="data",
269        label_paths=label_paths,
270        label_key="data",
271        patch_shape=patch_shape,
272        is_seg_dataset=True,
273        **kwargs
274    )
275
276
277def get_brats24_loader(
278    path: Union[os.PathLike, str],
279    batch_size: int,
280    patch_shape: Tuple[int, ...],
281    modality: Literal["t1n", "t1c", "t2w", "t2f"] = "t2f",
282    region: Optional[Literal[
283        "whole_tumor", "tumor_core", "enhancing_tissue",
284        "surrounding_flair_hyperintensity", "non_enhancing_tumor_core", "resection_cavity",
285    ]] = None,
286    resize_inputs: bool = False,
287    download: bool = False,
288    **kwargs
289) -> DataLoader:
290    """Get the BraTS 2024 post-treatment adult glioma dataloader for brain tumor segmentation.
291
292    Args:
293        path: Filepath to a folder where the data is downloaded for further processing.
294        batch_size: The batch size for training.
295        patch_shape: The patch shape to use for training.
296        modality: The MRI sequence. Either 't1n', 't1c', 't2w' or 't2f'.
297        region: The tumor region to use as a binary target, see `REGIONS`. If None, the sub-region ids are used.
298        resize_inputs: Whether to resize inputs to the desired patch shape.
299        download: Whether to download the data if it is not present.
300        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
301
302    Returns:
303        The DataLoader.
304    """
305    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
306    dataset = get_brats24_dataset(path, patch_shape, modality, region, resize_inputs, download, **ds_kwargs)
307    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)