torch_em.data.datasets.medical.cap_lv

  1"""The CAP LV dataset contains annotations for left ventricular myocardium segmentation in cine cardiac MRI.
  2
  3The data was curated for the LV Segmentation Challenge of the Cardiac Atlas Project
  4(https://www.cardiacatlas.org/lv-segmentation-challenge/), which was held at the STACOM 2011 workshop.
  5It consists of 200 cine MR studies of patients with myocardial infarction and impaired left ventricular
  6contraction, randomly drawn from the DETERMINE cohort, and split into 100 studies with reference
  7segmentations of the left ventricular myocardium and 100 studies without them. The images are
  8steady-state free precession acquisitions in a short-axis and a long-axis view, distributed as
  9de-identified DICOM, and the reference segmentations are binary mask images. Only the short-axis view is
 10used here, since the masks are defined on it.
 11
 12NOTE: The challenge calls the 100 annotated studies the 'test set' and the 100 studies without reference
 13segmentations the 'validation set', which is the opposite of the usual machine learning convention. This
 14module uses the annotated studies and ignores the others, so it does not expose a split argument.
 15
 16Unlike most cardiac benchmarks, the masks cover every frame of the cardiac cycle instead of only the end
 17diastole and the end systole phase. Each frame of the short-axis stack therefore becomes its own volume,
 18so that the 100 annotated studies amount to roughly 1600 annotated volumes.
 19
 20The labels are binary, see `LABEL_IDS`: 1 = left ventricular myocardium. The papillary muscles are
 21excluded from the myocardium.
 22
 23NOTE: The data is only available to users who signed the CAP Data Use Agreement, so it cannot be downloaded
 24automatically. To obtain the data, please follow these steps:
 25- Read the Terms and Conditions and the CAP Data Use Agreement linked at
 26  https://www.cardiacatlas.org/lv-segmentation-challenge/ and submit the request form at
 27  https://www.cardiacatlas.org/lv-segmentation-challenge/request-lv-segmentation-challenge/.
 28- After approval you receive login credentials for the download. Extract the downloaded archives into the
 29  folder passed as 'path', so that one folder per study (e.g. 'DET0000101') holds the DICOM images
 30  ('<study>_SA<slice>_ph<frame>.dcm', plus the long-axis images that are not used here) and the binary mask
 31  images of the annotated studies ('<study>_SA<slice>_ph<frame>.png').
 32
 33NOTE: The file naming convention is documented by the organizers. The study ids are the prefix 'DET' of the
 34DETERMINE cohort followed by 7 digits (https://www.cardiacatlas.org/determine/). The image filenames are of
 35the form 'DET0011201_SA3_ph0.dcm', see the annotation file page of the CAP LV Landmark Detection Challenge
 36(http://stacom.cardiacatlas.org/lv-landmark-detection-challenge/annotation-file/, which is distributed from
 37the same 200 studies), and the masks are PNG files with exactly the same stem, e.g.
 38'DET0000101_SA1_ph10.png', see the challenge FAQ (http://www.cardiacatlas.org/web/guest/faq). The slice and
 39the frame index are neither zero-padded nor 1-based.
 40
 41NOTE: The slice number in the filename is taken from the DICOM database and, as the challenge FAQ states
 42explicitly, does not always follow the anatomical slice order. The slices are therefore sorted by their
 43position along the slice normal, which is computed from the DICOM header, and not by the filename.
 44
 45NOTE: This module was written against the documented layout of the challenge distribution and could not be
 46validated on the data itself, because no openly published copy of it exists. It requires the pydicom
 47python package.
 48
 49The DICOM slices are stacked into volumes and stored in hdf5 files (the keys are 'raw' and 'labels').
 50
 51The data is only shared under the CAP Data Use Agreement for the DETERMINE cohort and the Terms and
 52Conditions of the consensus segmentation project, which restrict the use to that project and do not permit
 53redistribution, so please make sure that you are allowed to use the data for your purpose. NOTE: The CAP
 54states on https://www.cardiacatlas.org/determine/ that the DETERMINE data is currently unavailable for
 55download while the data sharing agreement with the data contributor is being renewed.
 56
 57This dataset is from the publication https://doi.org/10.1016/j.media.2013.09.001.
 58Please cite it if you use this dataset in your research.
 59"""
 60
 61import os
 62import re
 63from glob import glob
 64from tqdm import tqdm
 65from natsort import natsorted
 66from typing import Union, Tuple, List
 67
 68import numpy as np
 69
 70from torch.utils.data import Dataset, DataLoader
 71
 72import torch_em
 73
 74from .. import util
 75
 76
 77LABEL_IDS = {"background": 0, "myocardium": 1}
 78
 79SLICE_PATTERN = re.compile(r"^(?P<study>.+)_SA(?P<slice>\d+)_ph(?P<frame>\d+)$")
 80
 81
 82def _index_masks(path):
 83    """Index the binary masks by their stem, which is the same as the stem of the DICOM slice they belong to."""
 84    masks = {}
 85    for extension in ["png", "tif", "dcm"]:
 86        for mask_path in glob(os.path.join(path, "**", f"*_SA*_ph*.{extension}"), recursive=True):
 87            stem = os.path.splitext(os.path.basename(mask_path))[0]
 88            if extension == "dcm" and stem in masks:  # The images are DICOM as well, so PNG masks take precedence.
 89                continue
 90            masks.setdefault(stem, mask_path)
 91
 92    return masks
 93
 94
 95def _load_mask(mask_path):
 96    if mask_path.endswith(".dcm"):
 97        import pydicom
 98        return np.asarray(pydicom.dcmread(mask_path).pixel_array)
 99
100    import imageio.v3 as imageio
101    mask = np.asarray(imageio.imread(mask_path))
102    return mask[..., 0] if mask.ndim == 3 else mask
103
104
105def _slice_position(dcm):
106    """Compute the position of a slice along the slice normal, which orders the slices from apex to base."""
107    orientation = np.asarray(dcm.ImageOrientationPatient, dtype="float64")
108    normal = np.cross(orientation[:3], orientation[3:])
109    return float(np.dot(np.asarray(dcm.ImagePositionPatient, dtype="float64"), normal))
110
111
112def _group_slices(path):
113    """Group the short-axis DICOM slices of the studies into one volume per cardiac frame."""
114    volumes = {}
115    for image_path in natsorted(glob(os.path.join(path, "**", "*_SA*_ph*.dcm"), recursive=True)):
116        match = SLICE_PATTERN.match(os.path.splitext(os.path.basename(image_path))[0])
117        if match is None:
118            continue
119
120        key = (match.group("study"), int(match.group("frame")))
121        volumes.setdefault(key, []).append(image_path)
122
123    return volumes
124
125
126def _preprocess_inputs(path, preprocessed_dir):
127    import h5py
128    import pydicom
129
130    os.makedirs(preprocessed_dir, exist_ok=True)
131    volumes = _group_slices(path)
132    masks = _index_masks(path)
133
134    for (study, frame), image_paths in tqdm(sorted(volumes.items()), desc="Preprocessing the CAP LV studies"):
135        volume_path = os.path.join(preprocessed_dir, f"{study}_ph{frame:02}.h5")
136        if os.path.exists(volume_path):
137            continue
138
139        stems = [os.path.splitext(os.path.basename(p))[0] for p in image_paths]
140        if any(stem not in masks for stem in stems):  # The unannotated studies come without reference segmentations.
141            continue
142
143        # The slice number in the filename does not always follow the anatomical order, so the slices are
144        # sorted by their position along the slice normal instead, as recommended by the challenge FAQ.
145        slices = [(pydicom.dcmread(p), stem) for p, stem in zip(image_paths, stems)]
146        slices = sorted(slices, key=lambda item: _slice_position(item[0]))
147
148        raw = np.stack([np.asarray(dcm.pixel_array) for dcm, _ in slices])
149        labels = np.stack([_load_mask(masks[stem]) for _, stem in slices]) > 0
150
151        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
152        with h5py.File(f"{volume_path}.tmp", "w") as f:
153            f.create_dataset("raw", data=raw, compression="gzip")
154            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
155
156        os.rename(f"{volume_path}.tmp", volume_path)
157
158
159def get_cap_lv_data(path: Union[os.PathLike, str], download: bool = False) -> str:
160    """Obtain the CAP LV dataset.
161
162    Args:
163        path: Filepath to a folder where the manually downloaded data is stored.
164        download: Whether to download the data if it is not present. The data cannot be downloaded
165            automatically, so this raises if the data has not been downloaded manually.
166
167    Returns:
168        Filepath where the preprocessed data is stored.
169    """
170    preprocessed_dir = os.path.join(path, "preprocessed")
171    if glob(os.path.join(preprocessed_dir, "*.h5")):
172        return preprocessed_dir
173
174    if not glob(os.path.join(path, "**", "*_SA*_ph*.dcm"), recursive=True):
175        msg = "'torch_em' cannot download this dataset, because the LV Segmentation Challenge data is only "
176        msg += "available to users who signed the CAP Data Use Agreement. Please submit a request at "
177        msg += "'https://www.cardiacatlas.org/lv-segmentation-challenge/request-lv-segmentation-challenge/' and "
178        msg += f"extract the data you receive into '{path}', one folder per study."
179        raise NotImplementedError(msg)
180
181    _preprocess_inputs(path, preprocessed_dir)
182    return preprocessed_dir
183
184
185def get_cap_lv_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
186    """Get paths to the CAP LV data.
187
188    Args:
189        path: Filepath to a folder where the manually downloaded data is stored.
190        download: Whether to download the data if it is not present.
191
192    Returns:
193        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
194    """
195    data_dir = get_cap_lv_data(path, download)
196    return natsorted(glob(os.path.join(data_dir, "*.h5")))
197
198
199def get_cap_lv_dataset(
200    path: Union[os.PathLike, str],
201    patch_shape: Tuple[int, ...],
202    resize_inputs: bool = False,
203    download: bool = False,
204    **kwargs
205) -> Dataset:
206    """Get the CAP LV dataset for left ventricular myocardium segmentation.
207
208    Args:
209        path: Filepath to a folder where the manually downloaded data is stored.
210        patch_shape: The patch shape to use for training.
211        resize_inputs: Whether to resize inputs to the desired patch shape.
212        download: Whether to download the data if it is not present.
213        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
214
215    Returns:
216        The segmentation dataset.
217    """
218    volume_paths = get_cap_lv_paths(path, download)
219
220    if resize_inputs:
221        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
222        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
223            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
224        )
225
226    return torch_em.default_segmentation_dataset(
227        raw_paths=volume_paths,
228        raw_key="raw",
229        label_paths=volume_paths,
230        label_key="labels",
231        patch_shape=patch_shape,
232        is_seg_dataset=True,
233        **kwargs
234    )
235
236
237def get_cap_lv_loader(
238    path: Union[os.PathLike, str],
239    batch_size: int,
240    patch_shape: Tuple[int, ...],
241    resize_inputs: bool = False,
242    download: bool = False,
243    **kwargs
244) -> DataLoader:
245    """Get the CAP LV dataloader for left ventricular myocardium segmentation.
246
247    Args:
248        path: Filepath to a folder where the manually downloaded data is stored.
249        batch_size: The batch size for training.
250        patch_shape: The patch shape to use for training.
251        resize_inputs: Whether to resize inputs to the desired patch shape.
252        download: Whether to download the data if it is not present.
253        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
254
255    Returns:
256        The DataLoader.
257    """
258    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
259    dataset = get_cap_lv_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
260    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)