torch_em.data.datasets.medical.iugc2024

  1"""IUGC 2024 is the Intrapartum Ultrasound Grand Challenge dataset for fetal head and
  2pubic symphysis segmentation in transperineal ultrasound videos, recorded to assess the
  3progression of labor. Videos are stored per-frame and only a subset of frames per video
  4is annotated with a pixel-wise, 2-class (fetal head, pubic symphysis) segmentation mask.
  5
  6The dataset is located at
  7https://www.kaggle.com/datasets/aspirexxx/iugc-ultrasound-video-dataset-miccai-2024.
  8This dataset is from the publication https://doi.org/10.1007/978-3-031-96318-6_1.
  9Please cite it if you use this dataset for your research.
 10
 11NOTE: The dataset stores raw frames as videos (`.avi`) instead of individual images, so
 12this module extracts the specific annotated frame(s) out of each video and caches it to
 13disk as an image next to the corresponding mask. The "train" split stores masks nested
 14per-video (`seg/<video_stem>/mask/<video_stem>_<frame_index>_6.png`) and its video files
 15are named `<recording_id>__<video_stem>.avi`, unlike the flat layout of "val" and "test".
 16"""
 17
 18import os
 19import csv
 20from glob import glob
 21from typing import Union, Tuple, List, Literal
 22
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32KAGGLE_DATASET_NAME = "aspirexxx/iugc-ultrasound-video-dataset-miccai-2024"
 33
 34# The dataset does not expose a stable top-level folder layout, so the per-split prefixes
 35# below were resolved once via the Kaggle Files API (`KaggleApi.dataset_list_files`).
 36SPLIT_PREFIXES = {
 37    "train": "DatasetV3/train-20251119T060603Z-1-001/train/",
 38    "val": "DatasetV3/val-20251119T054616Z-1-001/val/",
 39    "test": "DatasetV3/test-20251119T054614Z-1-001/test/",
 40}
 41
 42
 43def _get_kaggle_api():
 44    try:
 45        from kaggle.api.kaggle_api_extended import KaggleApi
 46    except ModuleNotFoundError:
 47        msg = "Please install the Kaggle API. You can do this using 'pip install kaggle'. "
 48        msg += "After you have installed kaggle, you would need an API token. "
 49        msg += "Follow the instructions at https://www.kaggle.com/docs/api."
 50        raise ModuleNotFoundError(msg)
 51
 52    api = KaggleApi()
 53    api.authenticate()
 54    return api
 55
 56
 57def _list_seg_filenames(api, prefix):
 58    seg_prefix = f"{prefix}seg/"
 59    filenames = []
 60    token = None
 61    while True:
 62        response = api.dataset_list_files(KAGGLE_DATASET_NAME, page_token=token, page_size=500)
 63        names = [f.name for f in response.files]
 64        for name in names:
 65            if name.startswith(seg_prefix) and name.endswith(".png"):
 66                filenames.append(os.path.basename(name))
 67
 68        # The listing is alphabetically ordered, so once we have moved past the "seg/" folder
 69        # (and already collected some file names) we can stop early.
 70        if filenames and not any(name.startswith(seg_prefix) for name in names):
 71            break
 72
 73        token = response.next_page_token
 74        if not token:
 75            break
 76
 77    return filenames
 78
 79
 80def _list_all_filenames(api):
 81    filenames = []
 82    token = None
 83    while True:
 84        response = api.dataset_list_files(KAGGLE_DATASET_NAME, page_token=token, page_size=500)
 85        filenames.extend(f.name for f in response.files)
 86        token = response.next_page_token
 87        if not token:
 88            break
 89
 90    return filenames
 91
 92
 93def _match_train_video_name(video_names, video_stem):
 94    # Train videos are named "<recording_id>__<video_stem>.avi", unlike the flat
 95    # "<video_stem>.avi" naming used for the "val" and "test" splits.
 96    matches = [name for name in video_names if os.path.splitext(name)[0].rsplit("__", 1)[-1] == video_stem]
 97    if len(matches) != 1:
 98        raise RuntimeError(
 99            f"Found {len(matches)} candidate video files for '{video_stem}' in the train split, expected exactly 1."
100        )
101
102    return matches[0]
103
104
105def _download_train_data(api, prefix, rows, videos_dir, seg_dir):
106    all_filenames = _list_all_filenames(api)
107
108    videos_prefix = f"{prefix}videos/"
109    video_names = [name[len(videos_prefix):] for name in all_filenames if name.startswith(videos_prefix)]
110
111    seg_prefix = f"{prefix}seg/"
112    mask_names = [
113        name[len(seg_prefix):] for name in all_filenames
114        if name.startswith(seg_prefix) and name.endswith(".png") and "/mask/" in name
115    ]
116
117    for row in rows:
118        video_stem = os.path.splitext(row["filename"])[0]
119
120        video_name = _match_train_video_name(video_names, video_stem)
121        video_path = os.path.join(videos_dir, video_name)
122        if not os.path.exists(video_path):
123            api.dataset_download_file(KAGGLE_DATASET_NAME, f"{videos_prefix}{video_name}", path=videos_dir, quiet=False)
124
125        video_mask_names = [name for name in mask_names if name.startswith(f"{video_stem}/mask/")]
126        for mask_name in video_mask_names:
127            mask_path = os.path.join(seg_dir, os.path.basename(mask_name))
128            if os.path.exists(mask_path):
129                continue
130            api.dataset_download_file(KAGGLE_DATASET_NAME, f"{seg_prefix}{mask_name}", path=seg_dir, quiet=False)
131
132
133def get_iugc2024_data(
134    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
135) -> str:
136    """Download the IUGC 2024 dataset.
137
138    Args:
139        path: Filepath to a folder where the data is downloaded for further processing.
140        split: The choice of data split.
141        download: Whether to download the data if it is not present.
142
143    Returns:
144        Filepath where the data is downloaded.
145    """
146    if split not in SPLIT_PREFIXES:
147        raise ValueError(f"'{split}' is not a supported split. Choose one of {list(SPLIT_PREFIXES.keys())}.")
148
149    data_dir = os.path.join(path, split)
150    videos_dir = os.path.join(data_dir, "videos")
151    seg_dir = os.path.join(data_dir, "seg")
152    os.makedirs(videos_dir, exist_ok=True)
153    os.makedirs(seg_dir, exist_ok=True)
154
155    info_path = os.path.join(data_dir, "seg_info.csv")
156    if os.path.exists(info_path):
157        return data_dir
158
159    if not download:
160        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
161
162    prefix = SPLIT_PREFIXES[split]
163    api = _get_kaggle_api()
164
165    api.dataset_download_file(KAGGLE_DATASET_NAME, f"{prefix}seg/seg_info.csv", path=data_dir, quiet=False)
166
167    with open(info_path) as f:
168        rows = list(csv.DictReader(f))
169
170    if split == "train":
171        _download_train_data(api, prefix, rows, videos_dir, seg_dir)
172    else:
173        seg_filenames = _list_seg_filenames(api, prefix)
174        for row in rows:
175            video_name = row["filename"]
176            video_path = os.path.join(videos_dir, video_name)
177            if not os.path.exists(video_path):
178                api.dataset_download_file(
179                    KAGGLE_DATASET_NAME, f"{prefix}videos/{video_name}", path=videos_dir, quiet=False
180                )
181
182            video_stem = os.path.splitext(video_name)[0]
183            matches = [name for name in seg_filenames if name.startswith(video_stem)]
184            for mask_name in matches:
185                mask_path = os.path.join(seg_dir, mask_name)
186                if os.path.exists(mask_path):
187                    continue
188                api.dataset_download_file(KAGGLE_DATASET_NAME, f"{prefix}seg/{mask_name}", path=seg_dir, quiet=False)
189
190    return data_dir
191
192
193def get_iugc2024_paths(
194    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
195) -> Tuple[List[str], List[str]]:
196    """Get paths to the IUGC 2024 data.
197
198    Args:
199        path: Filepath to a folder where the data is downloaded for further processing.
200        split: The choice of data split.
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    import cv2
208
209    data_dir = get_iugc2024_data(path=path, split=split, download=download)
210    videos_dir = os.path.join(data_dir, "videos")
211    seg_dir = os.path.join(data_dir, "seg")
212
213    frames_dir = os.path.join(data_dir, "frames")
214    os.makedirs(frames_dir, exist_ok=True)
215
216    mask_paths = sorted(glob(os.path.join(seg_dir, "*.png")))
217
218    image_paths, gt_paths = [], []
219    for mask_path in mask_paths:
220        mask_stem = os.path.splitext(os.path.basename(mask_path))[0]
221        frame_path = os.path.join(frames_dir, f"{mask_stem}.tif")
222
223        image_paths.append(frame_path)
224        gt_paths.append(mask_path)
225        if os.path.exists(frame_path):
226            continue
227
228        if split == "train":
229            # The mask stem is "<video_stem>_<frame_index>_6" and the corresponding video is
230            # named "<recording_id>__<video_stem>.avi".
231            video_stem, frame_index, _ = mask_stem.rsplit("_", 2)
232            video_candidates = glob(os.path.join(videos_dir, f"*__{video_stem}.avi"))
233            if not video_candidates:
234                raise RuntimeError(f"Could not find a video file for the mask stem '{mask_stem}'.")
235            video_path = video_candidates[0]
236            frame_index = int(frame_index)
237        else:
238            # The mask stem is either "<video_stem>" (val split) or "<video_stem>_<frame_index>" (test split).
239            video_candidates = glob(os.path.join(videos_dir, f"{mask_stem}.avi"))
240            if video_candidates:
241                video_path = video_candidates[0]
242                frame_index = 0
243            else:
244                video_stem, frame_index = mask_stem.rsplit("_", 1)
245                video_path = os.path.join(videos_dir, f"{video_stem}.avi")
246                frame_index = int(frame_index)
247
248        capture = cv2.VideoCapture(video_path)
249        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
250        success, frame = capture.read()
251        capture.release()
252        if not success:
253            raise RuntimeError(f"Could not read frame {frame_index} from '{video_path}'.")
254
255        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
256        imageio.imwrite(frame_path, frame, compression="zlib")
257
258    return image_paths, gt_paths
259
260
261def get_iugc2024_dataset(
262    path: Union[os.PathLike, str],
263    patch_shape: Tuple[int, int],
264    split: Literal["train", "val", "test"],
265    resize_inputs: bool = False,
266    download: bool = False,
267    **kwargs
268) -> Dataset:
269    """Get the IUGC 2024 dataset for fetal head and pubic symphysis segmentation.
270
271    Args:
272        path: Filepath to a folder where the data is downloaded for further processing.
273        patch_shape: The patch shape to use for training.
274        split: The choice of data split.
275        resize_inputs: Whether to resize the inputs to the patch shape.
276        download: Whether to download the data if it is not present.
277        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
278
279    Returns:
280        The segmentation dataset.
281    """
282    image_paths, gt_paths = get_iugc2024_paths(path, split, download)
283
284    if resize_inputs:
285        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
286        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
287            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
288        )
289
290    return torch_em.default_segmentation_dataset(
291        raw_paths=image_paths,
292        raw_key=None,
293        label_paths=gt_paths,
294        label_key=None,
295        patch_shape=patch_shape,
296        is_seg_dataset=False,
297        **kwargs
298    )
299
300
301def get_iugc2024_loader(
302    path: Union[os.PathLike, str],
303    patch_shape: Tuple[int, int],
304    batch_size: int,
305    split: Literal["train", "val", "test"],
306    resize_inputs: bool = False,
307    download: bool = False,
308    **kwargs
309) -> DataLoader:
310    """Get the IUGC 2024 dataloader for fetal head and pubic symphysis segmentation.
311
312    Args:
313        path: Filepath to a folder where the data is downloaded for further processing.
314        patch_shape: The patch shape to use for training.
315        batch_size: The batch size for training.
316        split: The choice of data split.
317        resize_inputs: Whether to resize the inputs to the patch shape.
318        download: Whether to download the data if it is not present.
319        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
320
321    Returns:
322        The DataLoader.
323    """
324    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
325    dataset = get_iugc2024_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
326    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)