torch_em.data.datasets.medical.mvaa

  1"""The MVAA dataset contains annotations for mitral valve anatomy analysis across three imaging
  2modalities that reflect the clinical workflow from diagnosis to intervention: preoperative cardiac
  3CT (annular geometry), intraoperative 3D transesophageal echocardiography (3D TEE, leaflet
  4morphology) and surgical video (leaflet and instrument segmentation).
  5
  6The data was curated for the MVAA 2026 challenge (Mitral Valve Anatomy Analysis Using Multimodal
  7Imaging Data), held together with MICCAI 2026 (https://www.codabench.org/competitions/15662/). The
  8released training data consists of three genuinely separate distributions, exposed by this module
  9as three independent sets of functions:
 10- CT (`get_mvaa_ct_*`): 27 annotated volumes with a binary annulus mask, plus 1040 additional
 11  unlabeled volumes that are not exposed by this module.
 12- 3D TEE (`get_mvaa_tee_*`): 105 annotated ultrasound volumes with a 3-class label map (the exact
 13  semantics of the two foreground classes are not documented in the release).
 14- Surgical video (`get_mvaa_video_*`): 180 annotated RGB frames (30 frames each from 6 recordings),
 15  with polygon-derived instance masks for up to 17 anatomy and instrument classes, see
 16  `VIDEO_LABEL_IDS`. The mitral valve itself is class 10.
 17
 18NOTE: the validation split of each modality is released without ground truth (held out for the
 19official challenge evaluation), so this module only exposes the labeled training splits.
 20
 21The training data is distributed via Google Drive at
 22https://drive.google.com/file/d/14WneBUBZ1X4p69tRdRzximNb0IsWuh2B/view, as announced by the
 23organizers, see
 24https://communities.springernature.com/posts/miccai-2026-challenge-on-mitral-valve-multimodal-anatomical-analysis-challenge.
 25
 26This dataset is from https://doi.org/10.5281/zenodo.19726755.
 27Please cite it if you use this dataset in your research.
 28"""
 29
 30import os
 31import tarfile
 32from glob import glob
 33from tqdm import tqdm
 34from natsort import natsorted
 35from typing import Union, Tuple, List
 36
 37import numpy as np
 38
 39from torch.utils.data import Dataset, DataLoader
 40
 41import torch_em
 42
 43from .. import util
 44
 45
 46URL = "https://drive.google.com/uc?id=14WneBUBZ1X4p69tRdRzximNb0IsWuh2B"
 47CHECKSUM = "62695c24b269ae10962bd73bc6062ccb5a56bea0125e18325943b32f9e3a9bbf"
 48
 49CT_LABEL_IDS = {"background": 0, "annulus": 1}
 50
 51N_CT_VOLUMES = 27
 52N_TEE_VOLUMES = 105
 53N_VIDEO_FRAMES = 180
 54
 55VIDEO_LABEL_IDS = {
 56    "background": 0, "atrial_retractor": 1, "dissecting_forceps": 2, "scissors": 3, "needle_holder": 4,
 57    "sharp_knife": 5, "suture_organizer": 6, "suture": 7, "needle": 8, "atrial_inner_surface": 9,
 58    "mitral_valve": 10, "ventricle": 11, "blood": 12, "irrelevant": 13, "prosthetic_valve": 14,
 59    "annuloplasty_ring": 15, "gasket": 16, "valve_sizer": 17,
 60}
 61
 62
 63def get_mvaa_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 64    """Download the MVAA dataset.
 65
 66    Args:
 67        path: Filepath to a folder where the data is downloaded for further processing.
 68        download: Whether to download the data if it is not present.
 69
 70    Returns:
 71        Filepath to the folder with the raw 'reference_data' release.
 72    """
 73    data_dir = os.path.join(path, "reference_data")
 74    if os.path.exists(data_dir):
 75        return data_dir
 76
 77    os.makedirs(path, exist_ok=True)
 78    zip_path = os.path.join(path, "mvaa_train.zip")
 79    util.download_source_gdrive(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 80    util.unzip(zip_path=zip_path, dst=path)
 81
 82    return data_dir
 83
 84
 85def _preprocess_ct(data_dir, preprocessed_dir):
 86    import h5py
 87    import nibabel as nib
 88
 89    image_dir = os.path.join(data_dir, "t1_ct", "train", "labeled", "images")
 90    label_dir = os.path.join(data_dir, "t1_ct", "train", "labeled", "labels")
 91    image_paths = natsorted(glob(os.path.join(image_dir, "*.nii.gz")))
 92    os.makedirs(preprocessed_dir, exist_ok=True)
 93
 94    for image_path in tqdm(image_paths, desc="Preprocessing the MVAA CT volumes"):
 95        case_id = os.path.basename(image_path).split(".")[0]
 96        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 97        if os.path.exists(volume_path):
 98            continue
 99
100        label_path = os.path.join(label_dir, f"{case_id}-seg.nii.gz")
101
102        # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
103        raw = np.asarray(nib.load(image_path).dataobj).T
104        labels = np.asarray(nib.load(label_path).dataobj).T
105
106        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
107        with h5py.File(f"{volume_path}.tmp", "w") as f:
108            f.create_dataset("raw", data=raw, compression="gzip")
109            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
110
111        os.rename(f"{volume_path}.tmp", volume_path)
112
113
114def get_mvaa_ct_data(path: Union[os.PathLike, str], download: bool = False) -> str:
115    """Download and preprocess the MVAA cardiac CT data.
116
117    Args:
118        path: Filepath to a folder where the data is downloaded for further processing.
119        download: Whether to download the data if it is not present.
120
121    Returns:
122        Filepath where the preprocessed data is stored.
123    """
124    preprocessed_dir = os.path.join(path, "preprocessed", "ct")
125    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_CT_VOLUMES:
126        return preprocessed_dir
127
128    data_dir = get_mvaa_data(path, download)
129    _preprocess_ct(data_dir, preprocessed_dir)
130    return preprocessed_dir
131
132
133def get_mvaa_ct_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
134    """Get paths to the MVAA cardiac CT data.
135
136    Args:
137        path: Filepath to a folder where the data is downloaded for further processing.
138        download: Whether to download the data if it is not present.
139
140    Returns:
141        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
142    """
143    data_dir = get_mvaa_ct_data(path, download)
144    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
145    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{data_dir}'."
146    return volume_paths
147
148
149def get_mvaa_ct_dataset(
150    path: Union[os.PathLike, str],
151    patch_shape: Tuple[int, ...],
152    resize_inputs: bool = False,
153    download: bool = False,
154    **kwargs
155) -> Dataset:
156    """Get the MVAA dataset for mitral annulus segmentation in cardiac CT.
157
158    Args:
159        path: Filepath to a folder where the data is downloaded for further processing.
160        patch_shape: The patch shape to use for training.
161        resize_inputs: Whether to resize inputs to the desired patch shape.
162        download: Whether to download the data if it is not present.
163        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
164
165    Returns:
166        The segmentation dataset.
167    """
168    volume_paths = get_mvaa_ct_paths(path, download)
169
170    if resize_inputs:
171        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
172        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
173            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
174        )
175
176    return torch_em.default_segmentation_dataset(
177        raw_paths=volume_paths,
178        raw_key="raw",
179        label_paths=volume_paths,
180        label_key="labels",
181        patch_shape=patch_shape,
182        is_seg_dataset=True,
183        **kwargs
184    )
185
186
187def get_mvaa_ct_loader(
188    path: Union[os.PathLike, str],
189    batch_size: int,
190    patch_shape: Tuple[int, ...],
191    resize_inputs: bool = False,
192    download: bool = False,
193    **kwargs
194) -> DataLoader:
195    """Get the MVAA dataloader for mitral annulus segmentation in cardiac CT.
196
197    Args:
198        path: Filepath to a folder where the data is downloaded for further processing.
199        batch_size: The batch size for training.
200        patch_shape: The patch shape to use for training.
201        resize_inputs: Whether to resize inputs to the desired patch shape.
202        download: Whether to download the data if it is not present.
203        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
204
205    Returns:
206        The DataLoader.
207    """
208    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
209    dataset = get_mvaa_ct_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
210    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)
211
212
213def _preprocess_tee(data_dir, preprocessed_dir):
214    import h5py
215    import nibabel as nib
216
217    train_dir = os.path.join(data_dir, "t2_tee", "train")
218    image_paths = natsorted(glob(os.path.join(train_dir, "*-US.nii.gz")))
219    os.makedirs(preprocessed_dir, exist_ok=True)
220
221    for image_path in tqdm(image_paths, desc="Preprocessing the MVAA TEE volumes"):
222        case_id = os.path.basename(image_path).split("-US.nii.gz")[0]
223        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
224        if os.path.exists(volume_path):
225            continue
226
227        label_path = os.path.join(train_dir, f"{case_id}-label.nii.gz")
228
229        # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes.
230        raw = np.asarray(nib.load(image_path).dataobj).T
231        labels = np.asarray(nib.load(label_path).dataobj).T
232
233        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
234        with h5py.File(f"{volume_path}.tmp", "w") as f:
235            f.create_dataset("raw", data=raw, compression="gzip")
236            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
237
238        os.rename(f"{volume_path}.tmp", volume_path)
239
240
241def get_mvaa_tee_data(path: Union[os.PathLike, str], download: bool = False) -> str:
242    """Download and preprocess the MVAA 3D TEE data.
243
244    Args:
245        path: Filepath to a folder where the data is downloaded for further processing.
246        download: Whether to download the data if it is not present.
247
248    Returns:
249        Filepath where the preprocessed data is stored.
250    """
251    preprocessed_dir = os.path.join(path, "preprocessed", "tee")
252    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_TEE_VOLUMES:
253        return preprocessed_dir
254
255    data_dir = get_mvaa_data(path, download)
256    _preprocess_tee(data_dir, preprocessed_dir)
257    return preprocessed_dir
258
259
260def get_mvaa_tee_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
261    """Get paths to the MVAA 3D TEE data.
262
263    Args:
264        path: Filepath to a folder where the data is downloaded for further processing.
265        download: Whether to download the data if it is not present.
266
267    Returns:
268        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
269    """
270    data_dir = get_mvaa_tee_data(path, download)
271    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
272    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{data_dir}'."
273    return volume_paths
274
275
276def get_mvaa_tee_dataset(
277    path: Union[os.PathLike, str],
278    patch_shape: Tuple[int, ...],
279    resize_inputs: bool = False,
280    download: bool = False,
281    **kwargs
282) -> Dataset:
283    """Get the MVAA dataset for mitral valve leaflet segmentation in 3D TEE.
284
285    Args:
286        path: Filepath to a folder where the data is downloaded for further processing.
287        patch_shape: The patch shape to use for training.
288        resize_inputs: Whether to resize inputs to the desired 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`.
291
292    Returns:
293        The segmentation dataset.
294    """
295    volume_paths = get_mvaa_tee_paths(path, download)
296
297    if resize_inputs:
298        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
299        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
300            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
301        )
302
303    return torch_em.default_segmentation_dataset(
304        raw_paths=volume_paths,
305        raw_key="raw",
306        label_paths=volume_paths,
307        label_key="labels",
308        patch_shape=patch_shape,
309        is_seg_dataset=True,
310        **kwargs
311    )
312
313
314def get_mvaa_tee_loader(
315    path: Union[os.PathLike, str],
316    batch_size: int,
317    patch_shape: Tuple[int, ...],
318    resize_inputs: bool = False,
319    download: bool = False,
320    **kwargs
321) -> DataLoader:
322    """Get the MVAA dataloader for mitral valve leaflet segmentation in 3D TEE.
323
324    Args:
325        path: Filepath to a folder where the data is downloaded for further processing.
326        batch_size: The batch size for training.
327        patch_shape: The patch shape to use for training.
328        resize_inputs: Whether to resize inputs to the desired patch shape.
329        download: Whether to download the data if it is not present.
330        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
331
332    Returns:
333        The DataLoader.
334    """
335    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
336    dataset = get_mvaa_tee_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
337    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)
338
339
340def _preprocess_video(data_dir, preprocessed_dir):
341    import nibabel as nib
342    import imageio.v3 as imageio
343
344    image_dir = os.path.join(preprocessed_dir, "images")
345    label_dir = os.path.join(preprocessed_dir, "labels")
346    os.makedirs(image_dir, exist_ok=True)
347    os.makedirs(label_dir, exist_ok=True)
348
349    video_dirs = natsorted(glob(os.path.join(data_dir, "t3_vid", "train", "REC_*")))
350    for video_dir in tqdm(video_dirs, desc="Preprocessing the MVAA surgical video frames"):
351        tar_paths = natsorted(glob(os.path.join(video_dir, "*_png_Label.tar")))
352        for tar_path in tar_paths:
353            frame_name = os.path.basename(tar_path).split("_png_Label.tar")[0]
354            image_path = os.path.join(image_dir, f"{frame_name}.png")
355            label_path = os.path.join(label_dir, f"{frame_name}.tif")
356            if os.path.exists(image_path) and os.path.exists(label_path):
357                continue
358
359            with tarfile.open(tar_path) as tar:
360                member = next(m for m in tar.getmembers() if m.name.endswith(".nii.gz"))
361                tar.extract(member, path=video_dir, filter="data")
362                nii_path = os.path.join(video_dir, member.name)
363
364            # The transpose maps the nifti axis order (X, Y) to the (Y, X) order used by the RGB frame.
365            label = np.asarray(nib.load(nii_path).dataobj).T
366            os.remove(nii_path)
367
368            frame_path = os.path.join(video_dir, f"{frame_name}.png")
369            imageio.imwrite(image_path, imageio.imread(frame_path))
370            imageio.imwrite(label_path, label.astype("uint8"))
371
372
373def get_mvaa_video_data(path: Union[os.PathLike, str], download: bool = False) -> str:
374    """Download and preprocess the MVAA surgical video data.
375
376    Args:
377        path: Filepath to a folder where the data is downloaded for further processing.
378        download: Whether to download the data if it is not present.
379
380    Returns:
381        Filepath where the preprocessed data is stored.
382    """
383    preprocessed_dir = os.path.join(path, "preprocessed", "video")
384    if len(glob(os.path.join(preprocessed_dir, "images", "*.png"))) == N_VIDEO_FRAMES:
385        return preprocessed_dir
386
387    data_dir = get_mvaa_data(path, download)
388    _preprocess_video(data_dir, preprocessed_dir)
389    return preprocessed_dir
390
391
392def get_mvaa_video_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
393    """Get paths to the MVAA surgical video data.
394
395    Args:
396        path: Filepath to a folder where the data is downloaded for further processing.
397        download: Whether to download the data if it is not present.
398
399    Returns:
400        List of filepaths for the image data.
401        List of filepaths for the label data.
402    """
403    preprocessed_dir = get_mvaa_video_data(path, download)
404    image_paths = natsorted(glob(os.path.join(preprocessed_dir, "images", "*.png")))
405    label_paths = natsorted(glob(os.path.join(preprocessed_dir, "labels", "*.tif")))
406    assert image_paths and len(image_paths) == len(label_paths), \
407        f"The images and labels for '{preprocessed_dir}' do not match."
408    return image_paths, label_paths
409
410
411def get_mvaa_video_dataset(
412    path: Union[os.PathLike, str],
413    patch_shape: Tuple[int, int],
414    resize_inputs: bool = False,
415    download: bool = False,
416    **kwargs
417) -> Dataset:
418    """Get the MVAA dataset for mitral valve and instrument segmentation in surgical video.
419
420    Args:
421        path: Filepath to a folder where the data is downloaded for further processing.
422        patch_shape: The patch shape to use for training.
423        resize_inputs: Whether to resize inputs to the desired patch shape.
424        download: Whether to download the data if it is not present.
425        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
426
427    Returns:
428        The segmentation dataset.
429    """
430    image_paths, label_paths = get_mvaa_video_paths(path, download)
431
432    if resize_inputs:
433        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
434        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
435            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
436        )
437
438    return torch_em.default_segmentation_dataset(
439        raw_paths=image_paths,
440        raw_key=None,
441        label_paths=label_paths,
442        label_key=None,
443        is_seg_dataset=False,
444        patch_shape=patch_shape,
445        **kwargs
446    )
447
448
449def get_mvaa_video_loader(
450    path: Union[os.PathLike, str],
451    batch_size: int,
452    patch_shape: Tuple[int, int],
453    resize_inputs: bool = False,
454    download: bool = False,
455    **kwargs
456) -> DataLoader:
457    """Get the MVAA dataloader for mitral valve and instrument segmentation in surgical video.
458
459    Args:
460        path: Filepath to a folder where the data is downloaded for further processing.
461        batch_size: The batch size for training.
462        patch_shape: The patch shape to use for training.
463        resize_inputs: Whether to resize inputs to the desired patch shape.
464        download: Whether to download the data if it is not present.
465        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
466
467    Returns:
468        The DataLoader.
469    """
470    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
471    dataset = get_mvaa_video_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
472    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)