torch_em.data.datasets.medical.valdo

  1"""The VALDO dataset contains annotations for cerebral microbleed segmentation in brain MRI.
  2
  3The data was curated for Task 2 of the VALDO challenge ('Where is VALDO? VAscular Lesions DetectiOn and
  4segmentation', https://valdo.grand-challenge.org/Task2), which was held at MICCAI 2021 and targets the
  5detection and segmentation of small vessel disease markers. The public training release of Task 2 consists
  6of 72 subjects pooled from three cohorts (the subject ids 'sub-1xx', 'sub-2xx' and 'sub-3xx' correspond to
  711, 34 and 27 subjects respectively), so this module provides 72 annotated volumes. The test data of the
  8challenge was never released.
  9
 10Three co-registered sequences are available per subject and can be selected with the 'modality' argument:
 11the T2* scan ('t2s'), a T2-weighted scan ('t2') and a T1-weighted scan ('t1'). All of them are resampled to
 12the T2* space, in which the microbleeds are annotated, and the voxel size differs between the three cohorts.
 13NOTE: The T1 and T2 volumes have NaNs in the masked-out background, which are set to zero by this module.
 14
 15The annotations are binary, see `LABEL_IDS`: 1 = cerebral microbleed. The challenge evaluation derives
 16individual lesions from them with a connected component analysis with a neighbourhood of 6, so instance
 17labels can be obtained with `scipy.ndimage.label`.
 18
 19The three sequences and the annotations of a subject are bundled into one hdf5 file per subject by this
 20module, with the slice axis first (the keys are 'raw/t2s', 'raw/t2', 'raw/t1' and 'labels').
 21
 22The data is located at https://doi.org/10.5281/zenodo.4687995 and is licensed under CC BY-NC-SA 4.0.
 23
 24This dataset is from the publication https://doi.org/10.48550/arXiv.2208.07167.
 25Please cite it if you use this dataset in your research.
 26"""
 27
 28import os
 29from glob import glob
 30from tqdm import tqdm
 31from natsort import natsorted
 32from typing import Union, Tuple, List, Literal
 33
 34import numpy as np
 35
 36from torch.utils.data import Dataset, DataLoader
 37
 38import torch_em
 39
 40from .. import util
 41
 42
 43URL = "https://zenodo.org/records/4687995/files/Task2_v2.tar.gz"
 44CHECKSUM = "d47a9104deffd6a3a813ae53bf61209f734cce2818aa11611c1daf71188806b4"
 45
 46LABEL_IDS = {"background": 0, "microbleed": 1}
 47
 48MODALITIES = {"t2s": "T2S", "t2": "T2", "t1": "T1"}
 49
 50
 51def _preprocess_inputs(data_dir, preprocessed_dir):
 52    import h5py
 53    import nibabel as nib
 54
 55    subject_dirs = natsorted(glob(os.path.join(data_dir, "sub-*")))
 56    os.makedirs(preprocessed_dir, exist_ok=True)
 57
 58    for subject_dir in tqdm(subject_dirs, desc="Preprocessing the VALDO subjects"):
 59        subject_id = os.path.basename(subject_dir)
 60        volume_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
 61        if os.path.exists(volume_path):
 62            continue
 63
 64        # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
 65        label_path = os.path.join(subject_dir, f"{subject_id}_space-T2S_CMB.nii.gz")
 66        labels = np.asarray(nib.load(label_path).dataobj).T > 0
 67
 68        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
 69        with h5py.File(f"{volume_path}.tmp", "w") as f:
 70            for modality, suffix in MODALITIES.items():
 71                raw_path = os.path.join(subject_dir, f"{subject_id}_space-T2S_desc-masked_{suffix}.nii.gz")
 72                # The background of the T1 and T2 scans is NaN, which would propagate into the normalization.
 73                raw = np.nan_to_num(np.asarray(nib.load(raw_path).dataobj, dtype="float32").T)
 74                f.create_dataset(f"raw/{modality}", data=raw, compression="gzip")
 75
 76            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
 77
 78        os.rename(f"{volume_path}.tmp", volume_path)
 79
 80
 81def get_valdo_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 82    """Download the VALDO Task 2 dataset.
 83
 84    Args:
 85        path: Filepath to a folder where the data is downloaded for further processing.
 86        download: Whether to download the data if it is not present.
 87
 88    Returns:
 89        Filepath where the preprocessed data is stored.
 90    """
 91    preprocessed_dir = os.path.join(path, "preprocessed")
 92    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == 72:
 93        return preprocessed_dir
 94
 95    os.makedirs(path, exist_ok=True)
 96
 97    data_dir = os.path.join(path, "Task2")
 98    if not os.path.exists(data_dir):
 99        tar_path = os.path.join(path, "Task2_v2.tar.gz")
100        util.download_source(path=tar_path, url=URL, download=download, checksum=CHECKSUM)
101        util.unzip_tarfile(tar_path=tar_path, dst=path, remove=False)
102
103    _preprocess_inputs(data_dir, preprocessed_dir)
104    return preprocessed_dir
105
106
107def get_valdo_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
108    """Get paths to the VALDO Task 2 data.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        download: Whether to download the data if it is not present.
113
114    Returns:
115        List of filepaths for the hdf5 files, which contain the image data ('raw/<modality>')
116        and the label data ('labels').
117    """
118    data_dir = get_valdo_data(path, download)
119    volume_paths = natsorted(glob(os.path.join(data_dir, "sub-*.h5")))
120    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{data_dir}'."
121    return volume_paths
122
123
124def get_valdo_dataset(
125    path: Union[os.PathLike, str],
126    patch_shape: Tuple[int, ...],
127    modality: Literal["t2s", "t2", "t1"] = "t2s",
128    resize_inputs: bool = False,
129    download: bool = False,
130    **kwargs
131) -> Dataset:
132    """Get the VALDO Task 2 dataset for cerebral microbleed segmentation.
133
134    Args:
135        path: Filepath to a folder where the data is downloaded for further processing.
136        patch_shape: The patch shape to use for training.
137        modality: The MRI sequence. Either 't2s', 't2' or 't1'.
138        resize_inputs: Whether to resize inputs to the desired patch shape.
139        download: Whether to download the data if it is not present.
140        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
141
142    Returns:
143        The segmentation dataset.
144    """
145    if modality not in MODALITIES:
146        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {list(MODALITIES.keys())}.")
147
148    volume_paths = get_valdo_paths(path, download)
149
150    if resize_inputs:
151        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
152        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
153            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
154        )
155
156    return torch_em.default_segmentation_dataset(
157        raw_paths=volume_paths,
158        raw_key=f"raw/{modality}",
159        label_paths=volume_paths,
160        label_key="labels",
161        patch_shape=patch_shape,
162        is_seg_dataset=True,
163        **kwargs
164    )
165
166
167def get_valdo_loader(
168    path: Union[os.PathLike, str],
169    batch_size: int,
170    patch_shape: Tuple[int, ...],
171    modality: Literal["t2s", "t2", "t1"] = "t2s",
172    resize_inputs: bool = False,
173    download: bool = False,
174    **kwargs
175) -> DataLoader:
176    """Get the VALDO Task 2 dataloader for cerebral microbleed segmentation.
177
178    Args:
179        path: Filepath to a folder where the data is downloaded for further processing.
180        batch_size: The batch size for training.
181        patch_shape: The patch shape to use for training.
182        modality: The MRI sequence. Either 't2s', 't2' or 't1'.
183        resize_inputs: Whether to resize inputs to the desired patch shape.
184        download: Whether to download the data if it is not present.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    dataset = get_valdo_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
192    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)