torch_em.data.datasets.medical.maternal_fetal_us_video

  1"""The Maternal-Fetal Ultrasound Video dataset contains annotations for segmentation of the pubic
  2symphysis and the fetal head in transperineal intrapartum ultrasound videos.
  3
  4The dataset consists of 774 videos (68,106 frames in total) collected from three medical centers
  5(JNU, SYSU and SMU) with different ultrasound devices. A subset of the frames is manually annotated
  6with pixel-level semantic segmentation masks for two structures: the pubic symphysis (class 1) and
  7the fetal head (class 2). This wrapper extracts the annotated frames from the videos (using the frame
  8index encoded in the annotation filename) and pairs them with the corresponding segmentation mask, so
  9that the data can be used as a regular 2d image segmentation dataset.
 10
 11The dataset ships three splits with slightly different storage conventions:
 12- 'train': videos with annotated frames are stored in 'train/pos/<video_id>/', with the video itself
 13  ('<video_id>.avi') and the per-frame masks ('mask/<video_id>_<frame_idx>_*.png') in the same folder.
 14  The masks use pixel values 7 (pubic symphysis) and 8 (fetal head), which are remapped to 1 and 2
 15  respectively to match the 'test' and 'val' splits (see NOTE below).
 16- 'test': videos are stored in 'test/video/Source*_*/<video_id>.avi' and the per-frame masks in
 17  'test/seg_label/<video_id>_<frame_idx>.npy' (pixel values 0, 1, 2).
 18- 'val': videos are stored in 'val/videos/<video_id>_<frame_idx>.avi' (the frame index is part of the
 19  video filename) and the single annotated mask per video in 'val/label_seg/<video_id>_seg.npy'
 20  (pixel values 0, 1, 2).
 21
 22NOTE: The 'train' split masks use different pixel values (7, 8) than the 'test' and 'val' splits
 23(1, 2). This wrapper remaps the 'train' masks to the same (1, 2) convention when caching the frames,
 24so that all splits share one label convention: 0 (background), 1 (pubic symphysis), 2 (fetal head).
 25
 26This dataset is located at https://doi.org/10.5281/zenodo.16869288.
 27This dataset is from the publication https://doi.org/10.1038/s41597-026-06900-5.
 28The dataset is licensed under CC-BY-4.0.
 29Please cite the publication above if you use this dataset for your research.
 30"""
 31
 32import os
 33import re
 34from glob import glob
 35from tqdm import tqdm
 36from natsort import natsorted
 37from typing import Union, Tuple, List, Literal
 38
 39import numpy as np
 40import imageio.v3 as imageio
 41
 42from torch.utils.data import Dataset, DataLoader
 43
 44import torch_em
 45
 46from .. import util
 47
 48
 49URL = "https://zenodo.org/records/16869288/files/IUGC2024.rar"
 50CHECKSUM = "048730551419192aeaf7b23bcf6dbdcda26ab2d4acde4f5007b350128d188052"
 51
 52SPLITS = ["train", "val", "test"]
 53
 54# The 'train' split masks use 7 / 8 for the two foreground classes, the 'test' / 'val' splits use 1 / 2.
 55TRAIN_LABEL_MAP = {0: 0, 7: 1, 8: 2}
 56
 57
 58def get_maternal_fetal_us_video_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 59    """Download the Maternal-Fetal Ultrasound Video dataset.
 60
 61    Args:
 62        path: Filepath to a folder where the data is downloaded for further processing.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        Filepath where the data is downloaded.
 67    """
 68    data_dir = os.path.join(path, "IUGC2024")
 69    if os.path.exists(data_dir):
 70        return data_dir
 71
 72    os.makedirs(path, exist_ok=True)
 73
 74    rar_path = os.path.join(path, "IUGC2024.rar")
 75    util.download_source(path=rar_path, url=URL, download=download, checksum=CHECKSUM)
 76    util.unzip_rarfile(rar_path=rar_path, dst=path)
 77
 78    return data_dir
 79
 80
 81def _read_frame(video_path, frame_idx):
 82    import cv2
 83
 84    cap = cv2.VideoCapture(video_path)
 85    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
 86    success, frame = cap.read()
 87    cap.release()
 88    if not success:
 89        raise RuntimeError(f"Could not read frame {frame_idx} from '{video_path}'.")
 90    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 91
 92
 93def _cache_train_split(data_dir, cache_dir):
 94    video_dirs = natsorted(glob(os.path.join(data_dir, "train", "pos", "*")))
 95
 96    image_paths, label_paths = [], []
 97    for video_dir in tqdm(video_dirs, desc="Caching 'train' frames"):
 98        video_id = os.path.basename(video_dir)
 99        video_path = os.path.join(video_dir, f"{video_id}.avi")
100        mask_paths = natsorted(glob(os.path.join(video_dir, "mask", f"{video_id}_*.png")))
101
102        for mask_path in mask_paths:
103            fname = os.path.basename(mask_path)
104            rest = fname[len(video_id):-len(".png")].strip("_")
105            frame_idx = int(rest.split("_")[0])
106
107            image_path = os.path.join(cache_dir, f"train_{video_id}_{frame_idx}_image.tif")
108            label_path = os.path.join(cache_dir, f"train_{video_id}_{frame_idx}_label.tif")
109            image_paths.append(image_path)
110            label_paths.append(label_path)
111            if os.path.exists(image_path) and os.path.exists(label_path):
112                continue
113
114            frame = _read_frame(video_path, frame_idx)
115            mask = imageio.imread(mask_path)
116            label = np.vectorize(TRAIN_LABEL_MAP.get)(mask).astype("uint8")
117
118            imageio.imwrite(image_path, frame, compression="zlib")
119            imageio.imwrite(label_path, label, compression="zlib")
120
121    return image_paths, label_paths
122
123
124def _cache_test_split(data_dir, cache_dir):
125    mask_paths = natsorted(glob(os.path.join(data_dir, "test", "seg_label", "*.npy")))
126    video_paths = {
127        os.path.splitext(os.path.basename(p))[0]: p
128        for p in glob(os.path.join(data_dir, "test", "video", "Source*_*", "*.avi"))
129    }
130
131    image_paths, label_paths = [], []
132    for mask_path in tqdm(mask_paths, desc="Caching 'test' frames"):
133        fname = os.path.splitext(os.path.basename(mask_path))[0]
134        match = re.match(r"(.+)_(\d+)$", fname)
135        video_id, frame_idx = match.group(1), int(match.group(2))
136
137        video_path = video_paths.get(video_id)
138        if video_path is None:
139            raise RuntimeError(f"Could not find the video for '{video_id}' in the 'test' split.")
140
141        image_path = os.path.join(cache_dir, f"test_{video_id}_{frame_idx}_image.tif")
142        label_path = os.path.join(cache_dir, f"test_{video_id}_{frame_idx}_label.tif")
143        image_paths.append(image_path)
144        label_paths.append(label_path)
145        if os.path.exists(image_path) and os.path.exists(label_path):
146            continue
147
148        frame = _read_frame(video_path, frame_idx)
149        label = np.load(mask_path).astype("uint8")
150
151        imageio.imwrite(image_path, frame, compression="zlib")
152        imageio.imwrite(label_path, label, compression="zlib")
153
154    return image_paths, label_paths
155
156
157def _cache_val_split(data_dir, cache_dir):
158    mask_paths = natsorted(glob(os.path.join(data_dir, "val", "label_seg", "*_seg.npy")))
159    video_paths = {
160        re.match(r"(.+)_(\d+)$", os.path.splitext(os.path.basename(p))[0]).group(1): p
161        for p in glob(os.path.join(data_dir, "val", "videos", "*.avi"))
162    }
163
164    image_paths, label_paths = [], []
165    for mask_path in tqdm(mask_paths, desc="Caching 'val' frames"):
166        video_id = os.path.basename(mask_path)[:-len("_seg.npy")]
167
168        video_path = video_paths.get(video_id)
169        if video_path is None:
170            raise RuntimeError(f"Could not find the video for '{video_id}' in the 'val' split.")
171
172        frame_idx = int(re.match(r".+_(\d+)$", os.path.splitext(os.path.basename(video_path))[0]).group(1))
173
174        image_path = os.path.join(cache_dir, f"val_{video_id}_{frame_idx}_image.tif")
175        label_path = os.path.join(cache_dir, f"val_{video_id}_{frame_idx}_label.tif")
176        image_paths.append(image_path)
177        label_paths.append(label_path)
178        if os.path.exists(image_path) and os.path.exists(label_path):
179            continue
180
181        frame = _read_frame(video_path, frame_idx)
182        label = np.load(mask_path).astype("uint8")
183
184        imageio.imwrite(image_path, frame, compression="zlib")
185        imageio.imwrite(label_path, label, compression="zlib")
186
187    return image_paths, label_paths
188
189
190def get_maternal_fetal_us_video_paths(
191    path: Union[os.PathLike, str], split: Literal["train", "val", "test", "all"] = "all", download: bool = False,
192) -> Tuple[List[str], List[str]]:
193    """Get paths to the Maternal-Fetal Ultrasound Video data.
194
195    This extracts the annotated frames from the videos and caches them (together with their
196    segmentation masks) as tif files, so that repeated calls avoid re-decoding the videos.
197
198    Args:
199        path: Filepath to a folder where the data is downloaded for further processing.
200        split: The choice of data split. One of 'train', 'val', 'test' or 'all' (uses all splits).
201        download: Whether to download the data if it is not present.
202
203    Returns:
204        List of filepaths for the image data.
205        List of filepaths for the label data.
206    """
207    if split == "all":
208        splits = SPLITS
209    elif split in SPLITS:
210        splits = [split]
211    else:
212        raise ValueError(f"'{split}' is not a valid split. Choose from {SPLITS + ['all']}.")
213
214    data_dir = get_maternal_fetal_us_video_data(path, download)
215
216    cache_dir = os.path.join(path, "frames")
217    os.makedirs(cache_dir, exist_ok=True)
218
219    cache_funcs = {"train": _cache_train_split, "test": _cache_test_split, "val": _cache_val_split}
220
221    image_paths, label_paths = [], []
222    for this_split in splits:
223        this_image_paths, this_label_paths = cache_funcs[this_split](data_dir, cache_dir)
224        image_paths.extend(this_image_paths)
225        label_paths.extend(this_label_paths)
226
227    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
228
229    return image_paths, label_paths
230
231
232def get_maternal_fetal_us_video_dataset(
233    path: Union[os.PathLike, str],
234    patch_shape: Tuple[int, int],
235    split: Literal["train", "val", "test", "all"] = "all",
236    resize_inputs: bool = False,
237    download: bool = False,
238    **kwargs
239) -> Dataset:
240    """Get the Maternal-Fetal Ultrasound Video dataset for pubic symphysis and fetal head segmentation.
241
242    Args:
243        path: Filepath to a folder where the data is downloaded for further processing.
244        patch_shape: The patch shape to use for training.
245        split: The choice of data split. One of 'train', 'val', 'test' or 'all' (uses all splits).
246        resize_inputs: Whether to resize the inputs to the expected patch shape.
247        download: Whether to download the data if it is not present.
248        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
249
250    Returns:
251        The segmentation dataset.
252    """
253    image_paths, label_paths = get_maternal_fetal_us_video_paths(path, split, download)
254
255    if resize_inputs:
256        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
257        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
258            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
259        )
260
261    return torch_em.default_segmentation_dataset(
262        raw_paths=image_paths,
263        raw_key=None,
264        label_paths=label_paths,
265        label_key=None,
266        patch_shape=patch_shape,
267        is_seg_dataset=False,
268        **kwargs
269    )
270
271
272def get_maternal_fetal_us_video_loader(
273    path: Union[os.PathLike, str],
274    batch_size: int,
275    patch_shape: Tuple[int, int],
276    split: Literal["train", "val", "test", "all"] = "all",
277    resize_inputs: bool = False,
278    download: bool = False,
279    **kwargs
280) -> DataLoader:
281    """Get the Maternal-Fetal Ultrasound Video dataloader for pubic symphysis and fetal head segmentation.
282
283    Args:
284        path: Filepath to a folder where the data is downloaded for further processing.
285        batch_size: The batch size for training.
286        patch_shape: The patch shape to use for training.
287        split: The choice of data split. One of 'train', 'val', 'test' or 'all' (uses all splits).
288        resize_inputs: Whether to resize the inputs to the expected patch shape.
289        download: Whether to download the data if it is not present.
290        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
291
292    Returns:
293        The DataLoader.
294    """
295    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
296    dataset = get_maternal_fetal_us_video_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
297    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)