torch_em.data.datasets.medical.brats

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