torch_em.data.datasets.medical.emidec

  1"""The EMIDEC dataset contains annotations for the left ventricle and for myocardial infarction
  2in delayed-enhancement (late gadolinium enhanced) cardiac MRI.
  3
  4The data was curated for the EMIDEC challenge (https://emidec.com), which was held together with the STACOM
  5workshop at MICCAI 2020. The public release of the segmentation contest consists of 100 training cases,
  633 normal cases (case ids with the letter 'N') and 67 pathological cases (case ids with the letter 'P').
  7The 50 cases of the test set are distributed without ground truth and are therefore not exposed here,
  8so that this module provides 100 annotated volumes out of the 150 exams of the database.
  9
 10The label ids are described in `LABEL_IDS`: 0 = background, 1 = cavity, 2 = normal myocardium,
 113 = myocardial infarction, 4 = no-reflow (permanent microvascular obstruction). The whole myocardium is
 12the union of the ids 2, 3 and 4.
 13
 14The DE-MRI are stored as nifti volumes with the slice axis last, so they are converted to hdf5 volumes
 15with the slice axis first (the keys are 'raw' and 'labels') by this module.
 16
 17The data is located at https://emidec.com/dataset and is licensed under CC BY-NC-SA 4.0.
 18
 19This dataset is from the publication https://doi.org/10.3390/data5040089.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from typing import Union, Tuple, List, Optional, Literal
 28
 29import numpy as np
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38URL = "https://emidec.com/dataset/download"
 39CHECKSUM = "9270495a800d717092a6c21feba4b4f20fe90c61be85ac2a0cfc35570e321528"
 40
 41LABEL_IDS = {"background": 0, "cavity": 1, "myocardium": 2, "infarction": 3, "no_reflow": 4}
 42
 43PATHOLOGIES = {"normal": "N", "pathological": "P"}
 44
 45
 46def _preprocess_inputs(data_dir, preprocessed_dir):
 47    import h5py
 48    import nibabel as nib
 49
 50    case_dirs = [p for p in natsorted(glob(os.path.join(data_dir, "Case_*"))) if os.path.isdir(p)]
 51    os.makedirs(preprocessed_dir, exist_ok=True)
 52
 53    for case_dir in tqdm(case_dirs, desc="Preprocessing the EMIDEC cases"):
 54        case_id = os.path.basename(case_dir)
 55        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 56        if os.path.exists(volume_path):
 57            continue
 58
 59        # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
 60        raw = np.asarray(nib.load(os.path.join(case_dir, "Images", f"{case_id}.nii.gz")).dataobj).T
 61        labels = np.asarray(nib.load(os.path.join(case_dir, "Contours", f"{case_id}.nii.gz")).dataobj).T
 62
 63        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
 64        with h5py.File(f"{volume_path}.tmp", "w") as f:
 65            f.create_dataset("raw", data=raw, compression="gzip")
 66            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
 67
 68        os.rename(f"{volume_path}.tmp", volume_path)
 69
 70
 71def get_emidec_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 72    """Download the EMIDEC dataset.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        Filepath where the preprocessed data is stored.
 80    """
 81    preprocessed_dir = os.path.join(path, "preprocessed")
 82    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == 100:
 83        return preprocessed_dir
 84
 85    os.makedirs(path, exist_ok=True)
 86
 87    data_dir = os.path.join(path, "emidec-dataset-1.0.1")
 88    if not os.path.exists(data_dir):
 89        zip_path = os.path.join(path, "emidec-dataset-1.0.1.zip")
 90        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 91        util.unzip(zip_path=zip_path, dst=path)
 92
 93    _preprocess_inputs(data_dir, preprocessed_dir)
 94    return preprocessed_dir
 95
 96
 97def get_emidec_paths(
 98    path: Union[os.PathLike, str],
 99    pathology: Optional[Literal["normal", "pathological"]] = None,
100    download: bool = False,
101) -> List[str]:
102    """Get paths to the EMIDEC data.
103
104    Args:
105        path: Filepath to a folder where the data is downloaded for further processing.
106        pathology: The choice of cases. Either 'normal' or 'pathological'. If None, all cases are returned.
107        download: Whether to download the data if it is not present.
108
109    Returns:
110        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
111    """
112    data_dir = get_emidec_data(path, download)
113
114    if pathology is not None and pathology not in PATHOLOGIES:
115        raise ValueError(f"'{pathology}' is not a valid pathology. Please choose one of {list(PATHOLOGIES.keys())}.")
116
117    prefix = "*" if pathology is None else PATHOLOGIES[pathology]
118    volume_paths = natsorted(glob(os.path.join(data_dir, f"Case_{prefix}*.h5")))
119    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{data_dir}'."
120
121    return volume_paths
122
123
124def get_emidec_dataset(
125    path: Union[os.PathLike, str],
126    patch_shape: Tuple[int, ...],
127    pathology: Optional[Literal["normal", "pathological"]] = None,
128    resize_inputs: bool = False,
129    download: bool = False,
130    **kwargs
131) -> Dataset:
132    """Get the EMIDEC dataset for myocardial infarction 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        pathology: The choice of cases. Either 'normal' or 'pathological'. If None, all cases are returned.
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    volume_paths = get_emidec_paths(path, pathology, download)
146
147    if resize_inputs:
148        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
149        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
150            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
151        )
152
153    return torch_em.default_segmentation_dataset(
154        raw_paths=volume_paths,
155        raw_key="raw",
156        label_paths=volume_paths,
157        label_key="labels",
158        patch_shape=patch_shape,
159        is_seg_dataset=True,
160        **kwargs
161    )
162
163
164def get_emidec_loader(
165    path: Union[os.PathLike, str],
166    batch_size: int,
167    patch_shape: Tuple[int, ...],
168    pathology: Optional[Literal["normal", "pathological"]] = None,
169    resize_inputs: bool = False,
170    download: bool = False,
171    **kwargs
172) -> DataLoader:
173    """Get the EMIDEC dataloader for myocardial infarction segmentation.
174
175    Args:
176        path: Filepath to a folder where the data is downloaded for further processing.
177        batch_size: The batch size for training.
178        patch_shape: The patch shape to use for training.
179        pathology: The choice of cases. Either 'normal' or 'pathological'. If None, all cases are returned.
180        resize_inputs: Whether to resize inputs to the desired patch shape.
181        download: Whether to download the data if it is not present.
182        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
183
184    Returns:
185        The DataLoader.
186    """
187    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
188    dataset = get_emidec_dataset(path, patch_shape, pathology, resize_inputs, download, **ds_kwargs)
189    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)