torch_em.data.datasets.medical.atriaseg

  1"""The AtriaSeg dataset contains annotations for left atrium segmentation in
  2late gadolinium enhanced (LGE) cardiac MRI of patients with atrial fibrillation.
  3
  4The data was curated for the 2018 Atrial Segmentation Challenge
  5(https://www.cardiacatlas.org/atriaseg2018-challenge/), which was held together with the STACOM workshop at
  6MICCAI 2018. The release consists of 154 3D LGE-MRI of 154 patients, split into the 100 studies of the
  7official training set and the 54 studies of the official test set, which is selected with the 'split'
  8argument. All studies come with a segmentation of the left atrial cavity ('laendo.nrrd') and of the left
  9atrial wall ('lawall.nrrd'), which are merged into one label volume with the ids described in `LABEL_IDS`:
 101 = left atrium cavity, 2 = left atrium wall. The two structures are disjoint. The challenge itself only
 11evaluated the left atrium cavity.
 12
 13The LGE-MRI are stored as NRRD volumes with the slice axis last, so they are converted to hdf5 volumes with
 14the slice axis first (the keys are 'raw' and 'labels') by this module. All volumes have 88 slices.
 15
 16NOTE: This requires the pynrrd python package.
 17
 18The data is located at https://www.cardiacatlas.org/atriaseg2018-challenge/atria-seg-data/.
 19
 20This dataset is from the publication https://doi.org/10.1016/j.media.2020.101832.
 21Please cite it if you use this dataset in your research.
 22"""
 23
 24import os
 25from glob import glob
 26from tqdm import tqdm
 27from natsort import natsorted
 28from typing import Union, Tuple, List, Literal
 29
 30import numpy as np
 31
 32from torch.utils.data import Dataset, DataLoader
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39URL = "https://www.dropbox.com/scl/fi/nero2nlaocdcdfhzwu5h0/2018_UTAH_MICCAI.zip?rlkey=vkkfrkc2l6x1e61jqyutb35qn&dl=1"  # noqa
 40CHECKSUM = "bee5ee5bd19a1caa1a375e147e56e7e691a4bc64e3873dc672d9d2b963a8f5e0"
 41
 42LABEL_IDS = {"background": 0, "la_cavity": 1, "la_wall": 2}
 43
 44SPLITS = {"train": "Training Set", "test": "Testing Set"}
 45
 46N_VOLUMES = {"train": 100, "test": 54}
 47
 48
 49def _preprocess_inputs(split_dir, preprocessed_dir):
 50    import h5py
 51    import nrrd
 52
 53    case_dirs = [p for p in natsorted(glob(os.path.join(split_dir, "*"))) if os.path.isdir(p)]
 54    os.makedirs(preprocessed_dir, exist_ok=True)
 55
 56    for case_dir in tqdm(case_dirs, desc=f"Preprocessing the AtriaSeg cases of '{os.path.basename(split_dir)}'"):
 57        case_id = os.path.basename(case_dir)
 58        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 59        if os.path.exists(volume_path):
 60            continue
 61
 62        # The transpose maps the NRRD axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
 63        raw = nrrd.read(os.path.join(case_dir, "lgemri.nrrd"))[0].T
 64        cavity = nrrd.read(os.path.join(case_dir, "laendo.nrrd"))[0].T
 65        wall = nrrd.read(os.path.join(case_dir, "lawall.nrrd"))[0].T
 66
 67        # The masks are stored with the foreground value 255, which is mapped to the semantic label ids here.
 68        labels = np.zeros(raw.shape, dtype="uint8")
 69        labels[cavity > 0] = LABEL_IDS["la_cavity"]
 70        labels[wall > 0] = LABEL_IDS["la_wall"]
 71
 72        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
 73        with h5py.File(f"{volume_path}.tmp", "w") as f:
 74            f.create_dataset("raw", data=raw, compression="gzip")
 75            f.create_dataset("labels", data=labels, compression="gzip")
 76
 77        os.rename(f"{volume_path}.tmp", volume_path)
 78
 79
 80def get_atriaseg_data(
 81    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 82) -> str:
 83    """Download the AtriaSeg dataset.
 84
 85    Args:
 86        path: Filepath to a folder where the data is downloaded for further processing.
 87        split: The choice of data split. Either 'train' or 'test'.
 88        download: Whether to download the data if it is not present.
 89
 90    Returns:
 91        Filepath where the preprocessed data is stored.
 92    """
 93    if split not in SPLITS:
 94        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(SPLITS.keys())}.")
 95
 96    preprocessed_dir = os.path.join(path, "preprocessed", split)
 97    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_VOLUMES[split]:
 98        return preprocessed_dir
 99
100    os.makedirs(path, exist_ok=True)
101
102    split_dir = os.path.join(path, SPLITS[split])
103    if not os.path.exists(split_dir):
104        zip_path = os.path.join(path, "2018_UTAH_MICCAI.zip")
105        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
106        util.unzip(zip_path=zip_path, dst=path)
107
108    _preprocess_inputs(split_dir, preprocessed_dir)
109    return preprocessed_dir
110
111
112def get_atriaseg_paths(
113    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
114) -> List[str]:
115    """Get paths to the AtriaSeg data.
116
117    Args:
118        path: Filepath to a folder where the data is downloaded for further processing.
119        split: The choice of data split. Either 'train' or 'test'.
120        download: Whether to download the data if it is not present.
121
122    Returns:
123        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
124    """
125    data_dir = get_atriaseg_data(path, split, download)
126    return natsorted(glob(os.path.join(data_dir, "*.h5")))
127
128
129def get_atriaseg_dataset(
130    path: Union[os.PathLike, str],
131    patch_shape: Tuple[int, ...],
132    split: Literal["train", "test"],
133    resize_inputs: bool = False,
134    download: bool = False,
135    **kwargs
136) -> Dataset:
137    """Get the AtriaSeg dataset for left atrium segmentation.
138
139    Args:
140        path: Filepath to a folder where the data is downloaded for further processing.
141        patch_shape: The patch shape to use for training.
142        split: The choice of data split. Either 'train' or 'test'.
143        resize_inputs: Whether to resize inputs to the desired patch shape.
144        download: Whether to download the data if it is not present.
145        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
146
147    Returns:
148        The segmentation dataset.
149    """
150    volume_paths = get_atriaseg_paths(path, split, download)
151
152    if resize_inputs:
153        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
154        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
155            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
156        )
157
158    return torch_em.default_segmentation_dataset(
159        raw_paths=volume_paths,
160        raw_key="raw",
161        label_paths=volume_paths,
162        label_key="labels",
163        patch_shape=patch_shape,
164        is_seg_dataset=True,
165        **kwargs
166    )
167
168
169def get_atriaseg_loader(
170    path: Union[os.PathLike, str],
171    batch_size: int,
172    patch_shape: Tuple[int, ...],
173    split: Literal["train", "test"],
174    resize_inputs: bool = False,
175    download: bool = False,
176    **kwargs
177) -> DataLoader:
178    """Get the AtriaSeg dataloader for left atrium segmentation.
179
180    Args:
181        path: Filepath to a folder where the data is downloaded for further processing.
182        batch_size: The batch size for training.
183        patch_shape: The patch shape to use for training.
184        split: The choice of data split. Either 'train' or 'test'.
185        resize_inputs: Whether to resize inputs to the desired patch shape.
186        download: Whether to download the data if it is not present.
187        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
188
189    Returns:
190        The DataLoader.
191    """
192    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
193    dataset = get_atriaseg_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
194    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)