torch_em.data.datasets.medical.brainmetshare

  1"""The BrainMetShare dataset contains annotations for brain metastasis segmentation in multi-sequence brain MRI.
  2
  3It comprises 156 whole brain MRI studies (105 with radiologist-drawn metastasis segmentations, 51 unlabeled)
  4with four co-registered, skull-stripped 3D sequences resampled to 256 x 256 pixels: T1 gradient-echo
  5post-contrast, T1 spin-echo pre-contrast, T1 spin-echo post-contrast and T2 FLAIR post-contrast.
  6Only the 105 labeled studies are provided by this dataset.
  7
  8NOTE: The label legend is as follows:
  9- background: 0, metastasis: 1
 10Verified on the data: the label volumes only contain the ids 0 and 1.
 11
 12The data is a redistribution of the official Stanford release at
 13https://www.kaggle.com/datasets/kapilesha/brainmetshare-nii, which stores the four sequences and the
 14segmentation of each study as nifti files (uploaded under the MIT license; the underlying data is subject
 15to the Stanford University Dataset Research Use Agreement). The official release at
 16https://aimi.stanford.edu/brainmetshare requires registration, so please make sure that you are allowed
 17to use the data for your purpose.
 18
 19NOTE: The official release can also be used. Download it as described below and this dataset will use it
 20instead of the redistribution:
 21- Visit https://aimi.stanford.edu/brainmetshare and follow the link to the dataset on Stanford's Redivis
 22  platform (https://stanford.redivis.com/datasets/1emj-bjxt3p6s0).
 23- Register / log in, fill in your contact details and accept the research use agreement.
 24- Download the data (e.g. with 'azcopy' as described on the website) and place it such that the labeled cases
 25  are located at '<path>/mets_stanford_release_train/<case>/{0,1,2,3,seg}'. Each modality folder holds the
 26  slices of one sequence ('0': T1 gradient-echo post, '1': T1 spin-echo pre, '2': T1 spin-echo post,
 27  '3': T2 FLAIR post) and 'seg' holds the binary metastasis mask (0, 255).
 28
 29This dataset is from the publication https://doi.org/10.1002/jmri.26766.
 30Please cite it if you use this dataset in your research.
 31"""
 32
 33import os
 34from glob import glob
 35from tqdm import tqdm
 36from natsort import natsorted
 37from typing import Union, Tuple, Literal, List
 38
 39import numpy as np
 40
 41from torch.utils.data import Dataset, DataLoader
 42
 43import torch_em
 44
 45from .. import util
 46
 47
 48KAGGLE_DATASET = "kapilesha/brainmetshare-nii"
 49
 50LABEL_IDS = {"background": 0, "metastasis": 1}
 51
 52MODALITIES = {"t1_gre_post": "0", "t1_se_pre": "1", "t1_se_post": "2", "flair": "3"}
 53
 54# The names used for the sequences in the redistributed nifti files.
 55NIFTI_NAMES = {"t1_gre_post": "bravo", "t1_se_pre": "t1_pre", "t1_se_post": "t1_gd", "flair": "flair"}
 56
 57
 58def _load_volume(case_dir, nifti_name, folder_name):
 59    """Load a volume stored either as a nifti file or as a stack of 2d slice images."""
 60    nifti_path = os.path.join(case_dir, f"{nifti_name}.nii")
 61    if os.path.exists(nifti_path):
 62        import nibabel as nib
 63        return np.asarray(nib.load(nifti_path).dataobj).T  # (Z, Y, X)
 64
 65    import imageio.v3 as imageio
 66    folder = os.path.join(case_dir, folder_name)
 67    slice_paths = natsorted([p for p in glob(os.path.join(folder, "*")) if os.path.isfile(p)])
 68    assert len(slice_paths) > 0, f"Could not find any image files in {folder}."
 69    volume = np.stack([imageio.imread(p) for p in slice_paths])
 70    if volume.ndim == 4:  # Multi-channel slice images (e.g. RGB pngs) are reduced to a single channel.
 71        volume = volume[..., 0]
 72    return volume
 73
 74
 75def _preprocess_inputs(path, case_dirs):
 76    import h5py
 77
 78    preprocessed_dir = os.path.join(path, "preprocessed")
 79    os.makedirs(preprocessed_dir, exist_ok=True)
 80
 81    for case_dir in tqdm(case_dirs, desc="Preprocessing the BrainMetShare cases"):
 82        case_id = os.path.basename(case_dir)
 83        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 84        if os.path.exists(volume_path):
 85            continue
 86
 87        labels = (_load_volume(case_dir, "seg", "seg") > 0).astype(np.uint8)
 88
 89        # The file is written to a temporary path first, so that an interrupted run does not leave a corrupt file.
 90        with h5py.File(f"{volume_path}.tmp", "w") as f:
 91            f.create_dataset("labels", data=labels, compression="gzip")
 92            for modality in MODALITIES:
 93                raw = _load_volume(case_dir, NIFTI_NAMES[modality], MODALITIES[modality])
 94                assert raw.shape == labels.shape, f"Shape mismatch for {case_id}: {raw.shape} vs. {labels.shape}."
 95                f.create_dataset(f"raw/{modality}", data=raw, compression="gzip")
 96
 97        os.rename(f"{volume_path}.tmp", volume_path)
 98
 99    return preprocessed_dir
100
101
102def _find_case_dirs(path):
103    # The labeled cases of the official release.
104    case_dirs = glob(os.path.join(path, "mets_stanford_release_train", "*"))
105    if len(case_dirs) == 0:  # The labeled cases of the redistribution.
106        case_dirs = glob(os.path.join(path, "train", "Mets_*"))
107    return natsorted([p for p in case_dirs if os.path.isdir(p)])
108
109
110def get_brainmetshare_data(path: Union[os.PathLike, str], download: bool = False) -> str:
111    """Download the BrainMetShare dataset.
112
113    Args:
114        path: Filepath to a folder where the data is downloaded for further processing.
115        download: Whether to download the data if it is not present.
116
117    Returns:
118        Filepath where the data is preprocessed.
119    """
120    preprocessed_dir = os.path.join(path, "preprocessed")
121    if os.path.exists(preprocessed_dir) and len(glob(os.path.join(preprocessed_dir, "*.h5"))) > 0:
122        return preprocessed_dir
123
124    case_dirs = _find_case_dirs(path)
125    if len(case_dirs) == 0:
126        os.makedirs(path, exist_ok=True)
127        util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET, download=download)
128        util.unzip(zip_path=os.path.join(path, "brainmetshare-nii.zip"), dst=path)
129        case_dirs = _find_case_dirs(path)
130
131    assert len(case_dirs) > 0, f"Could not find any BrainMetShare cases at '{path}'."
132
133    return _preprocess_inputs(path, case_dirs)
134
135
136def get_brainmetshare_paths(
137    path: Union[os.PathLike, str], download: bool = False
138) -> Tuple[List[str], List[str]]:
139    """Get paths to the BrainMetShare data.
140
141    Args:
142        path: Filepath to a folder where the data is downloaded for further processing.
143        download: Whether to download the data if it is not present.
144
145    Returns:
146        List of filepaths for the image data.
147        List of filepaths for the label data.
148    """
149    data_dir = get_brainmetshare_data(path, download)
150    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
151    return volume_paths, volume_paths
152
153
154def get_brainmetshare_dataset(
155    path: Union[os.PathLike, str],
156    patch_shape: Tuple[int, ...],
157    modality: Literal["t1_gre_post", "t1_se_pre", "t1_se_post", "flair"] = "t1_gre_post",
158    resize_inputs: bool = False,
159    download: bool = False,
160    **kwargs
161) -> Dataset:
162    """Get the BrainMetShare dataset for brain metastasis segmentation.
163
164    Args:
165        path: Filepath to a folder where the data is downloaded for further processing.
166        patch_shape: The patch shape to use for training.
167        modality: The MRI sequence. One of 't1_gre_post', 't1_se_pre', 't1_se_post' or 'flair'.
168        resize_inputs: Whether to resize inputs to the desired patch shape.
169        download: Whether to download the data if it is not present.
170        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
171
172    Returns:
173        The segmentation dataset.
174    """
175    if modality not in MODALITIES:
176        raise ValueError(f"'{modality}' is not a valid modality. Choose one of {list(MODALITIES)}.")
177
178    raw_paths, label_paths = get_brainmetshare_paths(path, download)
179
180    if resize_inputs:
181        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
182        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
183            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
184        )
185
186    return torch_em.default_segmentation_dataset(
187        raw_paths=raw_paths,
188        raw_key=f"raw/{modality}",
189        label_paths=label_paths,
190        label_key="labels",
191        patch_shape=patch_shape,
192        is_seg_dataset=True,
193        **kwargs
194    )
195
196
197def get_brainmetshare_loader(
198    path: Union[os.PathLike, str],
199    batch_size: int,
200    patch_shape: Tuple[int, ...],
201    modality: Literal["t1_gre_post", "t1_se_pre", "t1_se_post", "flair"] = "t1_gre_post",
202    resize_inputs: bool = False,
203    download: bool = False,
204    **kwargs
205) -> DataLoader:
206    """Get the BrainMetShare dataloader for brain metastasis segmentation.
207
208    Args:
209        path: Filepath to a folder where the data is downloaded for further processing.
210        batch_size: The batch size for training.
211        patch_shape: The patch shape to use for training.
212        modality: The MRI sequence. One of 't1_gre_post', 't1_se_pre', 't1_se_post' or 'flair'.
213        resize_inputs: Whether to resize inputs to the desired patch shape.
214        download: Whether to download the data if it is not present.
215        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
216
217    Returns:
218        The DataLoader.
219    """
220    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
221    dataset = get_brainmetshare_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
222    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)