torch_em.data.datasets.medical.cetus

  1"""The CETUS dataset contains annotations for left ventricle segmentation in
  23D echocardiography of the heart.
  3
  4The data was curated for the CETUS challenge (Challenge on Endocardial Three-dimensional Ultrasound
  5Segmentation, https://www.creatis.insa-lyon.fr/Challenge/CETUS/), which was held at MICCAI 2014. The public
  6release consists of the 3D echocardiographic sequences of 45 patients. For each patient the end-diastolic (ED)
  7and the end-systolic (ES) frame are extracted and annotated, which gives 90 annotated volumes. The annotation
  8is a binary mask of the left ventricle lumen (the endocardial surface), see `LABEL_IDS`. The volumes of a
  9single phase can be selected with the 'phase' argument.
 10
 11The volumes are stored as nifti files with the slice axis last and the masks use the foreground value 255,
 12so they are converted to hdf5 volumes with the slice axis first (the keys are 'raw' and 'labels') and the
 13masks are binarized by this module. The image intensities are 8 bit and the voxels are isotropic.
 14
 15The data is located at https://humanheart-project.creatis.insa-lyon.fr/database/#collection/62eb991b73e9f0048c3a6c45
 16and is distributed under the CC BY-NC-SA 4.0 license.
 17
 18This dataset is from the publication https://doi.org/10.1109/TMI.2015.2503890.
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23from glob import glob
 24from tqdm import tqdm
 25from natsort import natsorted
 26from typing import Union, Tuple, List, Literal, Optional
 27
 28import numpy as np
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URL = "https://humanheart-project.creatis.insa-lyon.fr/database/api/v1/folder/62eb9a3e73e9f0048c3a6c46/download"
 38
 39# NOTE: The archive is created on the fly by the girder server, so its checksum changes with every download.
 40CHECKSUM = None
 41
 42LABEL_IDS = {"background": 0, "left_ventricle": 1}
 43
 44PHASES = ["ED", "ES"]
 45
 46N_VOLUMES = 90
 47
 48
 49def _preprocess_inputs(data_dir, preprocessed_dir):
 50    import h5py
 51    import nibabel as nib
 52
 53    case_dirs = natsorted(glob(os.path.join(data_dir, "patient*")))
 54    os.makedirs(preprocessed_dir, exist_ok=True)
 55
 56    for case_dir in tqdm(case_dirs, desc="Preprocessing the CETUS volumes"):
 57        case_id = os.path.basename(case_dir)
 58        for phase in PHASES:
 59            volume_path = os.path.join(preprocessed_dir, f"{case_id}_{phase}.h5")
 60            if os.path.exists(volume_path):
 61                continue
 62
 63            # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
 64            raw = np.asarray(nib.load(os.path.join(case_dir, f"{case_id}_{phase}.nii.gz")).dataobj).T
 65            mask = np.asarray(nib.load(os.path.join(case_dir, f"{case_id}_{phase}_gt.nii.gz")).dataobj).T
 66
 67            # The intensities are 8 bit values stored as floats and the mask uses the foreground value 255.
 68            raw = np.round(raw).astype("uint8")
 69            labels = (mask > 0).astype("uint8") * LABEL_IDS["left_ventricle"]
 70
 71            # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
 72            with h5py.File(f"{volume_path}.tmp", "w") as f:
 73                f.create_dataset("raw", data=raw, compression="gzip")
 74                f.create_dataset("labels", data=labels, compression="gzip")
 75
 76            os.rename(f"{volume_path}.tmp", volume_path)
 77
 78
 79def get_cetus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 80    """Download the CETUS dataset.
 81
 82    Args:
 83        path: Filepath to a folder where the data is downloaded for further processing.
 84        download: Whether to download the data if it is not present.
 85
 86    Returns:
 87        Filepath where the preprocessed data is stored.
 88    """
 89    preprocessed_dir = os.path.join(path, "preprocessed")
 90    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_VOLUMES:
 91        return preprocessed_dir
 92
 93    os.makedirs(path, exist_ok=True)
 94
 95    data_dir = os.path.join(path, "dataset")
 96    if not os.path.exists(data_dir):
 97        zip_path = os.path.join(path, "CETUS.zip")
 98        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 99        util.unzip(zip_path=zip_path, dst=path)
100
101    _preprocess_inputs(data_dir, preprocessed_dir)
102    return preprocessed_dir
103
104
105def get_cetus_paths(
106    path: Union[os.PathLike, str], phase: Optional[Literal["ED", "ES"]] = None, download: bool = False
107) -> List[str]:
108    """Get paths to the CETUS data.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        phase: The choice of cardiac phase. Either 'ED' or 'ES'. By default both phases are used.
113        download: Whether to download the data if it is not present.
114
115    Returns:
116        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
117    """
118    if phase is not None and phase not in PHASES:
119        raise ValueError(f"'{phase}' is not a valid cardiac phase. Please choose one of {PHASES}.")
120
121    data_dir = get_cetus_data(path, download)
122    volume_paths = natsorted(glob(os.path.join(data_dir, f"patient*_{'*' if phase is None else phase}.h5")))
123    assert len(volume_paths) > 0
124
125    return volume_paths
126
127
128def get_cetus_dataset(
129    path: Union[os.PathLike, str],
130    patch_shape: Tuple[int, ...],
131    phase: Optional[Literal["ED", "ES"]] = None,
132    resize_inputs: bool = False,
133    download: bool = False,
134    **kwargs
135) -> Dataset:
136    """Get the CETUS dataset for left ventricle segmentation.
137
138    Args:
139        path: Filepath to a folder where the data is downloaded for further processing.
140        patch_shape: The patch shape to use for training.
141        phase: The choice of cardiac phase. Either 'ED' or 'ES'. By default both phases are used.
142        resize_inputs: Whether to resize inputs to the desired patch shape.
143        download: Whether to download the data if it is not present.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
145
146    Returns:
147        The segmentation dataset.
148    """
149    volume_paths = get_cetus_paths(path, phase, download)
150
151    if resize_inputs:
152        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
153        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
154            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
155        )
156
157    return torch_em.default_segmentation_dataset(
158        raw_paths=volume_paths,
159        raw_key="raw",
160        label_paths=volume_paths,
161        label_key="labels",
162        patch_shape=patch_shape,
163        is_seg_dataset=True,
164        **kwargs
165    )
166
167
168def get_cetus_loader(
169    path: Union[os.PathLike, str],
170    batch_size: int,
171    patch_shape: Tuple[int, ...],
172    phase: Optional[Literal["ED", "ES"]] = None,
173    resize_inputs: bool = False,
174    download: bool = False,
175    **kwargs
176) -> DataLoader:
177    """Get the CETUS dataloader for left ventricle segmentation.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        batch_size: The batch size for training.
182        patch_shape: The patch shape to use for training.
183        phase: The choice of cardiac phase. Either 'ED' or 'ES'. By default both phases are used.
184        resize_inputs: Whether to resize inputs to the desired patch shape.
185        download: Whether to download the data if it is not present.
186        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
187
188    Returns:
189        The DataLoader.
190    """
191    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
192    dataset = get_cetus_dataset(path, patch_shape, phase, resize_inputs, download, **ds_kwargs)
193    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)