torch_em.data.datasets.medical.echonet_dynamic

  1"""EchoNet-Dynamic contains annotations for left ventricle segmentation in apical-4-chamber
  2echocardiography videos, together with cardiac function labels (ejection fraction, end-systolic
  3and end-diastolic volume).
  4
  5The dataset comprises 10,030 videos, collected at Stanford Health Care. Each video only has expert
  6tracings for the two labeled frames (end-diastole and end-systole) used to compute the left
  7ventricular ejection fraction, not dense per-frame masks. The dataset is located at
  8https://echonet.github.io/dynamic/ and distributed through the Stanford AIMI Center Shared Datasets
  9Portal under a non-commercial research use agreement: registration as an individual user is
 10required, re-distribution (including sharing the download link) is forbidden, and re-identification
 11attempts are prohibited. This module cannot download the dataset automatically; see
 12`get_echonet_dynamic_data` for the manual steps.
 13
 14This dataset is from the publication https://doi.org/10.1038/s41586-020-2145-8 (cite the DOI
 15https://doi.org/10.71718/yqp5-y078 for the data itself). Please cite them if you use this dataset in
 16your research.
 17
 18NOTE: Reading the videos requires 'opencv-python' ('cv2'), and the tracing rasterization requires
 19'scikit-image'.
 20"""
 21
 22import os
 23from glob import glob
 24from tqdm import tqdm
 25from natsort import natsorted
 26from typing import Union, Tuple, Literal, List
 27
 28import numpy as np
 29import pandas as pd
 30import imageio.v3 as imageio
 31
 32from torch.utils.data import Dataset, DataLoader
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39SPLITS = ("TRAIN", "VAL", "TEST")
 40
 41
 42def _trace_to_mask(x1, y1, x2, y2, shape):
 43    from skimage.draw import polygon
 44
 45    # Follows the rasterization used by the EchoNet dataset releases: the first coordinate pair is
 46    # the long axis of the left ventricle, and the remaining pairs are the perpendicular short-axis
 47    # distances. Walking the two sides of the short-axis points (one side forwards, the other
 48    # backwards) traces the outline of the traced region, which is then filled in.
 49    x = np.concatenate((x1[1:], x2[1:][::-1]))
 50    y = np.concatenate((y1[1:], y2[1:][::-1]))
 51
 52    rows, cols = polygon(np.round(y).astype(int), np.round(x).astype(int), shape=shape)
 53    mask = np.zeros(shape, dtype="uint8")
 54    mask[rows, cols] = 1
 55    return mask
 56
 57
 58def _find_video_path(videos_dir, file_name):
 59    stem = os.path.splitext(file_name)[0]
 60    for candidate in (file_name, f"{stem}.avi", f"{stem}.mp4"):
 61        candidate_path = os.path.join(videos_dir, candidate)
 62        if os.path.exists(candidate_path):
 63            return candidate_path
 64    return None
 65
 66
 67def _preprocess_inputs(data_dir, preprocessed_dir):
 68    import cv2
 69
 70    file_list = pd.read_csv(os.path.join(data_dir, "FileList.csv"))
 71    tracings = pd.read_csv(os.path.join(data_dir, "VolumeTracings.csv"))
 72
 73    videos_dir = os.path.join(data_dir, "Videos")
 74    os.makedirs(preprocessed_dir, exist_ok=True)
 75
 76    image_paths, gt_paths = [], []
 77    for _, row in tqdm(file_list.iterrows(), total=len(file_list), desc="Preprocess EchoNet-Dynamic"):
 78        file_name = str(row["FileName"])
 79        stem = os.path.splitext(file_name)[0]
 80        split = str(row["Split"]).upper()
 81
 82        video_tracings = tracings[tracings["FileName"] == file_name]
 83        if video_tracings.empty:
 84            video_tracings = tracings[tracings["FileName"] == stem + ".avi"]
 85        if video_tracings.empty:
 86            continue
 87
 88        video_path = _find_video_path(videos_dir, file_name)
 89        if video_path is None:
 90            continue
 91
 92        capture = cv2.VideoCapture(video_path)
 93
 94        for frame_idx, frame_tracings in video_tracings.groupby("Frame"):
 95            image_path = os.path.join(preprocessed_dir, f"{stem}_{split}_{frame_idx}.tif")
 96            mask_path = os.path.join(preprocessed_dir, f"{stem}_{split}_{frame_idx}_mask.tif")
 97
 98            if os.path.exists(image_path) and os.path.exists(mask_path):
 99                image_paths.append(image_path)
100                gt_paths.append(mask_path)
101                continue
102
103            capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_idx))
104            success, frame = capture.read()
105            if not success:
106                continue
107
108            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
109            mask = _trace_to_mask(
110                frame_tracings["X1"].to_numpy(),
111                frame_tracings["Y1"].to_numpy(),
112                frame_tracings["X2"].to_numpy(),
113                frame_tracings["Y2"].to_numpy(),
114                frame.shape,
115            )
116
117            imageio.imwrite(image_path, frame)
118            imageio.imwrite(mask_path, mask)
119
120            image_paths.append(image_path)
121            gt_paths.append(mask_path)
122
123        capture.release()
124
125    return image_paths, gt_paths
126
127
128def get_echonet_dynamic_data(path: Union[os.PathLike, str]) -> str:
129    """Obtain the EchoNet-Dynamic dataset.
130
131    NOTE: 'torch_em' cannot download this dataset, as it requires individual registration and
132    agreement to the EchoNet-Dynamic Research Use Agreement. Please follow these steps:
133    - Visit https://echonet.github.io/dynamic/ and follow the access instructions.
134    - Register (individually, per user) with the Stanford AIMI Center Shared Datasets Portal
135      (https://stanford.redivis.com/datasets/66s1-2hsmzj5rn) and agree to the Research Use Agreement
136      (non-commercial research use only, no re-distribution).
137    - Download the dataset and place it at `path`, so that it has the following structure:
138      `path/Videos`, `path/FileList.csv`, and `path/VolumeTracings.csv`.
139
140    Args:
141        path: Filepath to a folder where the dataset is stored.
142
143    Returns:
144        Filepath to the folder where the dataset is stored.
145    """
146    if not (
147        os.path.exists(os.path.join(path, "FileList.csv"))
148        and os.path.exists(os.path.join(path, "VolumeTracings.csv"))
149        and os.path.exists(os.path.join(path, "Videos"))
150    ):
151        raise RuntimeError(
152            f"Cannot find the EchoNet-Dynamic data at '{path}'. "
153            "This dataset requires manual download, see `get_echonet_dynamic_data` for the steps."
154        )
155
156    return path
157
158
159def get_echonet_dynamic_paths(
160    path: Union[os.PathLike, str], split: Literal["TRAIN", "VAL", "TEST", None] = None,
161) -> Tuple[List[str], List[str]]:
162    """Get paths to the EchoNet-Dynamic data.
163
164    Args:
165        path: Filepath to a folder where the dataset is stored.
166        split: The choice of data split. Either 'TRAIN', 'VAL' or 'TEST'. By default, all splits are used.
167
168    Returns:
169        List of filepaths for the image data.
170        List of filepaths for the label data.
171    """
172    data_dir = get_echonet_dynamic_data(path)
173
174    preprocessed_dir = os.path.join(path, "preprocessed")
175
176    image_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.tif")))
177    image_paths = [p for p in image_paths if not p.endswith("_mask.tif")]
178    if not image_paths:
179        image_paths, _ = _preprocess_inputs(data_dir, preprocessed_dir)
180        image_paths = natsorted(image_paths)
181
182    gt_paths = natsorted([p.replace(".tif", "_mask.tif") for p in image_paths])
183    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0
184
185    if split is not None:
186        assert split in SPLITS, f"'{split}' is not a valid split choice for the EchoNet-Dynamic dataset."
187        image_paths = [p for p in image_paths if f"_{split}_" in os.path.basename(p)]
188        gt_paths = [p for p in gt_paths if f"_{split}_" in os.path.basename(p)]
189
190    return image_paths, gt_paths
191
192
193def get_echonet_dynamic_dataset(
194    path: Union[os.PathLike, str],
195    patch_shape: Tuple[int, int],
196    split: Literal["TRAIN", "VAL", "TEST", None] = None,
197    resize_inputs: bool = False,
198    **kwargs
199) -> Dataset:
200    """Get the EchoNet-Dynamic dataset for left ventricle segmentation in echocardiography videos.
201
202    Args:
203        path: Filepath to a folder where the dataset is stored.
204        patch_shape: The patch shape to use for training.
205        split: The choice of data split. Either 'TRAIN', 'VAL' or 'TEST'. By default, all splits are used.
206        resize_inputs: Whether to resize inputs to the desired patch shape.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
208
209    Returns:
210        The segmentation dataset.
211    """
212    image_paths, gt_paths = get_echonet_dynamic_paths(path, split)
213
214    if resize_inputs:
215        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
216        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
217            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
218        )
219
220    return torch_em.default_segmentation_dataset(
221        raw_paths=image_paths,
222        raw_key=None,
223        label_paths=gt_paths,
224        label_key=None,
225        patch_shape=patch_shape,
226        is_seg_dataset=False,
227        **kwargs
228    )
229
230
231def get_echonet_dynamic_loader(
232    path: Union[os.PathLike, str],
233    batch_size: int,
234    patch_shape: Tuple[int, int],
235    split: Literal["TRAIN", "VAL", "TEST", None] = None,
236    resize_inputs: bool = False,
237    **kwargs
238) -> DataLoader:
239    """Get the EchoNet-Dynamic dataloader for left ventricle segmentation in echocardiography videos.
240
241    Args:
242        path: Filepath to a folder where the dataset is stored.
243        batch_size: The batch size for training.
244        patch_shape: The patch shape to use for training.
245        split: The choice of data split. Either 'TRAIN', 'VAL' or 'TEST'. By default, all splits are used.
246        resize_inputs: Whether to resize inputs to the desired patch shape.
247        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
248
249    Returns:
250        The DataLoader.
251    """
252    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
253    dataset = get_echonet_dynamic_dataset(path, patch_shape, split, resize_inputs, **ds_kwargs)
254    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)