torch_em.data.datasets.medical.isbi_mslesion

  1"""The ISBI MS Lesion dataset contains annotations for multiple sclerosis lesion segmentation
  2in longitudinal brain MRI.
  3
  4The data was curated for the Longitudinal Multiple Sclerosis Lesion Segmentation Challenge
  5(https://iacl.ece.jhu.edu/index.php/MSChallenge), which was held at ISBI 2015. The public training release
  6consists of 5 subjects that were scanned at 4 to 5 time points each, which amounts to 21 time points, and
  7each time point was delineated independently by 2 raters. This module therefore provides 21 volumes with
  82 sets of annotations, i.e. 42 annotated volumes, which are exposed together with the 'both' choice of the
  9'rater' argument. The 14 test subjects (61 further time points) are distributed without annotations and are
 10therefore not exposed here.
 11
 12Four co-registered sequences are available per time point and can be selected with the 'modality' argument:
 13a T2 FLAIR scan ('flair'), a T1-weighted MPRAGE scan ('mprage'), a proton density scan ('pd') and a
 14T2-weighted scan ('t2'). This module uses the preprocessed version of the scans (inhomogeneity corrected,
 15skull stripped and rigidly registered to a 1 mm isotropic MNI template), since only it is aligned with the
 16lesion masks. The scans in native acquisition space are also part of the download, but are not used here.
 17
 18The annotations are binary, see `LABEL_IDS`: 1 = multiple sclerosis lesion.
 19
 20The four sequences and the two sets of annotations of a time point are bundled into one hdf5 file per time
 21point by this module, with the slice axis first (the keys are 'raw/flair', 'raw/mprage', 'raw/pd', 'raw/t2',
 22'labels/rater1' and 'labels/rater2').
 23
 24The data is located at https://iacl.ece.jhu.edu/index.php/MSChallenge/data and may only be used for research
 25and education, see the license that is distributed with it.
 26
 27This dataset is from the publication https://doi.org/10.1016/j.neuroimage.2016.12.064.
 28Please cite it if you use this dataset in your research.
 29"""
 30
 31import os
 32from glob import glob
 33from tqdm import tqdm
 34from natsort import natsorted
 35from typing import Union, Tuple, List, Literal
 36
 37import numpy as np
 38
 39from torch.utils.data import Dataset, DataLoader
 40
 41import torch_em
 42
 43from ... import ConcatDataset
 44from .. import util
 45
 46
 47URL = "https://iacl.ece.jhu.edu/~aaron/data/training_final_v4.zip"
 48CHECKSUM = "f5db0c71d0dd90a19b156f8ab7eb6081dc51571c78c06646a7200bbb2b090be6"
 49
 50LABEL_IDS = {"background": 0, "ms_lesion": 1}
 51
 52MODALITIES = ["flair", "mprage", "pd", "t2"]
 53
 54RATERS = [1, 2]
 55
 56# The number of time points per training subject. They add up to the 21 time points of the training release.
 57TIME_POINTS = {"training01": 4, "training02": 4, "training03": 5, "training04": 4, "training05": 4}
 58
 59
 60def _preprocess_inputs(data_dir, preprocessed_dir):
 61    import h5py
 62    import nibabel as nib
 63
 64    os.makedirs(preprocessed_dir, exist_ok=True)
 65    for subject_id, n_time_points in tqdm(TIME_POINTS.items(), desc="Preprocessing the ISBI MS Lesion subjects"):
 66        for time_point in range(1, n_time_points + 1):
 67            case_id = f"{subject_id}_{time_point:02}"
 68            volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 69            if os.path.exists(volume_path):
 70                continue
 71
 72            # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
 73            # The file is written to a temporary path first, so an interrupted run leaves no corrupt file.
 74            with h5py.File(f"{volume_path}.tmp", "w") as f:
 75                for modality in MODALITIES:
 76                    raw_path = os.path.join(data_dir, subject_id, "preprocessed", f"{case_id}_{modality}_pp.nii")
 77                    raw = np.asarray(nib.load(raw_path).dataobj).T
 78                    f.create_dataset(f"raw/{modality}", data=raw, compression="gzip")
 79
 80                for rater in RATERS:
 81                    label_path = os.path.join(data_dir, subject_id, "masks", f"{case_id}_mask{rater}.nii")
 82                    labels = np.asarray(nib.load(label_path).dataobj).T > 0
 83                    f.create_dataset(f"labels/rater{rater}", data=labels.astype("uint8"), compression="gzip")
 84
 85            os.rename(f"{volume_path}.tmp", volume_path)
 86
 87
 88def get_isbi_mslesion_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 89    """Download the ISBI MS Lesion dataset.
 90
 91    Args:
 92        path: Filepath to a folder where the data is downloaded for further processing.
 93        download: Whether to download the data if it is not present.
 94
 95    Returns:
 96        Filepath where the preprocessed data is stored.
 97    """
 98    preprocessed_dir = os.path.join(path, "preprocessed")
 99    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == sum(TIME_POINTS.values()):
100        return preprocessed_dir
101
102    os.makedirs(path, exist_ok=True)
103
104    data_dir = os.path.join(path, "training")
105    if not os.path.exists(data_dir):
106        zip_path = os.path.join(path, "training_final_v4.zip")
107        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
108        util.unzip(zip_path=zip_path, dst=path, remove=False)
109
110    _preprocess_inputs(data_dir, preprocessed_dir)
111    return preprocessed_dir
112
113
114def get_isbi_mslesion_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
115    """Get paths to the ISBI MS Lesion data.
116
117    Args:
118        path: Filepath to a folder where the data is downloaded for further processing.
119        download: Whether to download the data if it is not present.
120
121    Returns:
122        List of filepaths for the hdf5 files, which contain the image data ('raw/<modality>')
123        and the label data ('labels/rater<rater>').
124    """
125    data_dir = get_isbi_mslesion_data(path, download)
126    volume_paths = natsorted(glob(os.path.join(data_dir, "training*.h5")))
127    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{data_dir}'."
128    return volume_paths
129
130
131def get_isbi_mslesion_dataset(
132    path: Union[os.PathLike, str],
133    patch_shape: Tuple[int, ...],
134    modality: Literal["flair", "mprage", "pd", "t2"] = "flair",
135    rater: Literal[1, 2, "both"] = "both",
136    resize_inputs: bool = False,
137    download: bool = False,
138    **kwargs
139) -> Dataset:
140    """Get the ISBI MS Lesion dataset for multiple sclerosis lesion segmentation.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        patch_shape: The patch shape to use for training.
145        modality: The MRI sequence. Either 'flair', 'mprage', 'pd' or 't2'.
146        rater: The choice of annotator. Either 1, 2 or 'both', which returns the volumes of both raters.
147        resize_inputs: Whether to resize inputs to the desired patch shape.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
150
151    Returns:
152        The segmentation dataset.
153    """
154    if modality not in MODALITIES:
155        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {MODALITIES}.")
156
157    if rater not in RATERS and rater != "both":
158        raise ValueError(f"'{rater}' is not a valid rater. Please choose one of {RATERS} or 'both'.")
159
160    volume_paths = get_isbi_mslesion_paths(path, download)
161
162    if resize_inputs:
163        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
164        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
165            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
166        )
167
168    datasets = []
169    for curr_rater in (RATERS if rater == "both" else [rater]):
170        datasets.append(
171            torch_em.default_segmentation_dataset(
172                raw_paths=volume_paths,
173                raw_key=f"raw/{modality}",
174                label_paths=volume_paths,
175                label_key=f"labels/rater{curr_rater}",
176                patch_shape=patch_shape,
177                is_seg_dataset=True,
178                **kwargs
179            )
180        )
181
182    return datasets[0] if len(datasets) == 1 else ConcatDataset(*datasets)
183
184
185def get_isbi_mslesion_loader(
186    path: Union[os.PathLike, str],
187    batch_size: int,
188    patch_shape: Tuple[int, ...],
189    modality: Literal["flair", "mprage", "pd", "t2"] = "flair",
190    rater: Literal[1, 2, "both"] = "both",
191    resize_inputs: bool = False,
192    download: bool = False,
193    **kwargs
194) -> DataLoader:
195    """Get the ISBI MS Lesion dataloader for multiple sclerosis lesion segmentation.
196
197    Args:
198        path: Filepath to a folder where the data is downloaded for further processing.
199        batch_size: The batch size for training.
200        patch_shape: The patch shape to use for training.
201        modality: The MRI sequence. Either 'flair', 'mprage', 'pd' or 't2'.
202        rater: The choice of annotator. Either 1, 2 or 'both', which returns the volumes of both raters.
203        resize_inputs: Whether to resize inputs to the desired patch shape.
204        download: Whether to download the data if it is not present.
205        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
206
207    Returns:
208        The DataLoader.
209    """
210    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
211    dataset = get_isbi_mslesion_dataset(path, patch_shape, modality, rater, resize_inputs, download, **ds_kwargs)
212    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)