torch_em.data.datasets.medical.upenn_gbm

  1"""The UPENN-GBM dataset contains annotations for brain tumor segmentation in multi-parametric MRI
  2of patients with de novo glioblastoma.
  3
  4It consists of 671 skull-stripped and co-registered MRI scans (630 patients, 240 x 240 x 155 voxels at 1 mm
  5isotropic resolution in SRI24 atlas space) with the modalities T1, T1GD (post-contrast T1), T2 and FLAIR.
  6The tumor sub-regions are labeled following the BraTS convention: 1: necrotic and non-enhancing tumor core,
  72: peritumoral edema, 4: GD-enhancing tumor.
  8Two kinds of segmentations are available: 147 scans have segmentations that were manually corrected and approved
  9by expert neuroradiologists ('manual', the default) and 611 scans have automatically generated segmentations
 10(label fusion of an ensemble of BraTS models, 'automated', which includes the 147 manually corrected scans).
 11
 12The dataset is located at https://www.cancerimagingarchive.net/collection/upenn-gbm/ and the nifti files are only
 13offered via IBM Aspera there. We download the data from a mirror of the original nifti files at
 14https://huggingface.co/datasets/MedOtter/UPENN-GBM instead.
 15
 16This dataset is from the publication https://doi.org/10.1038/s41597-022-01560-7.
 17The data was released at https://doi.org/10.7937/TCIA.709X-DN49.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import json
 23from tqdm import tqdm
 24from typing import Union, Tuple, Literal, List, Optional
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL = "https://huggingface.co/datasets/MedOtter/UPENN-GBM/resolve/main/"
 34
 35# The manifest lists all scans with their image and segmentation files. The nifti files are not checksummed.
 36CHECKSUM = "2b2c93a4181e9fc8e7db7d7b977cfb1981f63d13dcf19adc501df5cf92a70c76"
 37
 38MODALITIES = ["T1", "T1GD", "T2", "FLAIR"]
 39
 40LABEL_IDS = {"necrotic_core": 1, "edema": 2, "enhancing_tumor": 4}
 41
 42
 43def _get_scans(path, segmentation, download):
 44    manifest_path = os.path.join(path, "subjects_manifest.json")
 45    util.download_source(path=manifest_path, url=f"{URL}subjects_manifest.json", download=download, checksum=CHECKSUM)
 46    with open(manifest_path, "r") as f:
 47        scans = json.load(f)["subjects"]
 48    return [scan for scan in scans if scan[f"{segmentation}_segm"] is not None]
 49
 50
 51def get_upenn_gbm_data(
 52    path: Union[os.PathLike, str], segmentation: Literal["manual", "automated"] = "manual", download: bool = False
 53) -> str:
 54    """Download the UPENN-GBM dataset.
 55
 56    Args:
 57        path: Filepath to a folder where the data is downloaded for further processing.
 58        segmentation: The kind of segmentation. Either 'manual' (corrected by experts) or 'automated'.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        Filepath where the data is downloaded.
 63    """
 64    if segmentation not in ("manual", "automated"):
 65        raise ValueError(f"'{segmentation}' is not a valid segmentation. Please choose 'manual' or 'automated'.")
 66
 67    os.makedirs(path, exist_ok=True)
 68    scans = _get_scans(path, segmentation, download)
 69
 70    for scan in tqdm(scans, desc=f"Download UPENN-GBM ({segmentation} segmentations)"):
 71        for rel_path in list(scan["modalities"].values()) + [scan[f"{segmentation}_segm"]]:
 72            fpath = os.path.join(path, rel_path)
 73            os.makedirs(os.path.dirname(fpath), exist_ok=True)
 74            util.download_source(path=fpath, url=f"{URL}{rel_path}", download=download, checksum=None)
 75
 76    return path
 77
 78
 79def get_upenn_gbm_paths(
 80    path: Union[os.PathLike, str],
 81    modality: Optional[Literal["T1", "T1GD", "T2", "FLAIR"]] = None,
 82    segmentation: Literal["manual", "automated"] = "manual",
 83    download: bool = False,
 84) -> Tuple[List[Union[str, Tuple[str, ...]]], List[str]]:
 85    """Get paths to the UPENN-GBM data.
 86
 87    Args:
 88        path: Filepath to a folder where the data is downloaded for further processing.
 89        modality: The choice of modality. One of 'T1', 'T1GD', 'T2' or 'FLAIR'.
 90            By default, all modalities are returned as channels.
 91        segmentation: The kind of segmentation. Either 'manual' (corrected by experts) or 'automated'.
 92        download: Whether to download the data if it is not present.
 93
 94    Returns:
 95        List of filepaths for the image data.
 96        List of filepaths for the label data.
 97    """
 98    if modality is not None and modality not in MODALITIES:
 99        raise ValueError(f"'{modality}' is not a valid modality. Please choose from {MODALITIES}.")
100
101    data_dir = get_upenn_gbm_data(path, segmentation, download)
102    scans = _get_scans(data_dir, segmentation, download)
103
104    if modality is None:
105        raw_paths = [tuple(os.path.join(data_dir, scan["modalities"][mod]) for mod in MODALITIES) for scan in scans]
106    else:
107        raw_paths = [os.path.join(data_dir, scan["modalities"][modality]) for scan in scans]
108    label_paths = [os.path.join(data_dir, scan[f"{segmentation}_segm"]) for scan in scans]
109
110    return raw_paths, label_paths
111
112
113def get_upenn_gbm_dataset(
114    path: Union[os.PathLike, str],
115    patch_shape: Tuple[int, ...],
116    modality: Optional[Literal["T1", "T1GD", "T2", "FLAIR"]] = None,
117    segmentation: Literal["manual", "automated"] = "manual",
118    resize_inputs: bool = False,
119    download: bool = False,
120    **kwargs
121) -> Dataset:
122    """Get the UPENN-GBM dataset for brain tumor segmentation.
123
124    Args:
125        path: Filepath to a folder where the data is downloaded for further processing.
126        patch_shape: The patch shape to use for training.
127        modality: The choice of modality. One of 'T1', 'T1GD', 'T2' or 'FLAIR'.
128            By default, all modalities are used as channels.
129        segmentation: The kind of segmentation. Either 'manual' (corrected by experts) or 'automated'.
130        resize_inputs: Whether to resize inputs to the desired patch shape.
131        download: Whether to download the data if it is not present.
132        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
133
134    Returns:
135        The segmentation dataset.
136    """
137    raw_paths, label_paths = get_upenn_gbm_paths(path, modality, segmentation, download)
138
139    if resize_inputs:
140        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
141        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
142            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
143        )
144
145    return torch_em.default_segmentation_dataset(
146        raw_paths=raw_paths,
147        raw_key="data",
148        label_paths=label_paths,
149        label_key="data",
150        patch_shape=patch_shape,
151        is_seg_dataset=True,
152        with_channels=modality is None,
153        **kwargs
154    )
155
156
157def get_upenn_gbm_loader(
158    path: Union[os.PathLike, str],
159    batch_size: int,
160    patch_shape: Tuple[int, ...],
161    modality: Optional[Literal["T1", "T1GD", "T2", "FLAIR"]] = None,
162    segmentation: Literal["manual", "automated"] = "manual",
163    resize_inputs: bool = False,
164    download: bool = False,
165    **kwargs
166) -> DataLoader:
167    """Get the UPENN-GBM dataloader for brain tumor segmentation.
168
169    Args:
170        path: Filepath to a folder where the data is downloaded for further processing.
171        batch_size: The batch size for training.
172        patch_shape: The patch shape to use for training.
173        modality: The choice of modality. One of 'T1', 'T1GD', 'T2' or 'FLAIR'.
174            By default, all modalities are used as channels.
175        segmentation: The kind of segmentation. Either 'manual' (corrected by experts) or 'automated'.
176        resize_inputs: Whether to resize inputs to the desired patch shape.
177        download: Whether to download the data if it is not present.
178        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
179
180    Returns:
181        The DataLoader.
182    """
183    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
184    dataset = get_upenn_gbm_dataset(path, patch_shape, modality, segmentation, resize_inputs, download, **ds_kwargs)
185    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)