torch_em.data.datasets.medical.bagls_vf

  1"""BAGLS-VF extends the BAGLS glottis segmentation benchmark with pixel-wise annotations
  2of the left vocal fold, right vocal fold and glottal area (3 foreground classes), derived
  3from high-speed videolaryngoscopy recordings. Masks are RGB-coded, with a distinct color
  4per anatomical structure (encoded here as separate label values).
  5
  6The dataset is located at https://doi.org/10.5281/zenodo.19593658.
  7Please cite it if you use this dataset for your research.
  8"""
  9
 10import os
 11from glob import glob
 12from pathlib import Path
 13from typing import Union, Tuple, List, Literal
 14
 15import numpy as np
 16import imageio.v3 as imageio
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URLS = {
 26    "train": "https://zenodo.org/records/19593658/files/BAGLS-VF_train.zip",
 27    "test": "https://zenodo.org/records/19593658/files/BAGLS-VF_test.zip",
 28}
 29
 30CHECKSUMS = {
 31    "train": "accc5bc0acf920f344df5efb5b159fd15ad1a0941536542f65080a9f91b94c90",
 32    "test": "0876ae6461d5142b08338c4e45b9c96a4aa7642bebab4ce8d4a8508585b0907f",
 33}
 34
 35LABEL_MAP = {
 36    (0, 0, 0): 0,  # background
 37    (0, 0, 255): 1,  # vocal fold / glottis structure (blue)
 38    (0, 255, 0): 2,  # vocal fold / glottis structure (green)
 39    (255, 0, 0): 3,  # vocal fold / glottis structure (red)
 40}
 41
 42
 43def get_bagls_vf_data(path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False) -> str:
 44    """Download the BAGLS-VF dataset.
 45
 46    Args:
 47        path: Filepath to a folder where the data is downloaded for further processing.
 48        split: The choice of data split.
 49        download: Whether to download the data if it is not present.
 50
 51    Returns:
 52        Filepath where the data is downloaded.
 53    """
 54    data_dir = os.path.join(path, f"BAGLS-VF_{split}")
 55    if os.path.exists(data_dir):
 56        return data_dir
 57
 58    os.makedirs(path, exist_ok=True)
 59
 60    zip_path = os.path.join(path, f"BAGLS-VF_{split}.zip")
 61    util.download_source(path=zip_path, url=URLS[split], download=download, checksum=CHECKSUMS[split])
 62    util.unzip(zip_path=zip_path, dst=path)
 63
 64    return data_dir
 65
 66
 67def get_bagls_vf_paths(
 68    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 69) -> Tuple[List[str], List[str]]:
 70    """Get paths to the BAGLS-VF data.
 71
 72    Args:
 73        path: Filepath to a folder where the data is downloaded for further processing.
 74        split: The choice of data split.
 75        download: Whether to download the data if it is not present.
 76
 77    Returns:
 78        List of filepaths for the image data.
 79        List of filepaths for the label data.
 80    """
 81    data_dir = get_bagls_vf_data(path=path, split=split, download=download)
 82
 83    image_paths = sorted(glob(os.path.join(data_dir, f"{split}_images", "*.png")))
 84    mask_dir = os.path.join(data_dir, f"{split}_masks")
 85
 86    neu_gt_dir = os.path.join(data_dir, "preprocessed_masks")
 87    os.makedirs(neu_gt_dir, exist_ok=True)
 88
 89    gt_paths = []
 90    for image_path in image_paths:
 91        fname = Path(image_path).stem
 92        mask_path = os.path.join(mask_dir, f"{fname}.png")
 93        neu_gt_path = os.path.join(neu_gt_dir, f"{fname}.tif")
 94        gt_paths.append(neu_gt_path)
 95
 96        if os.path.exists(neu_gt_path):
 97            continue
 98
 99        mask = imageio.imread(mask_path)[..., :3]
100        instances = np.zeros(mask.shape[:2], dtype="uint8")
101        for color, label in LABEL_MAP.items():
102            if label == 0:
103                continue
104            binary_map = (mask == color).all(axis=-1)
105            instances[binary_map] = label
106
107        imageio.imwrite(neu_gt_path, instances, compression="zlib")
108
109    return image_paths, gt_paths
110
111
112def get_bagls_vf_dataset(
113    path: Union[os.PathLike, str],
114    patch_shape: Tuple[int, int],
115    split: Literal["train", "test"],
116    resize_inputs: bool = False,
117    download: bool = False,
118    **kwargs
119) -> Dataset:
120    """Get the BAGLS-VF dataset for vocal fold and glottal area segmentation.
121
122    Args:
123        path: Filepath to a folder where the data is downloaded for further processing.
124        patch_shape: The patch shape to use for training.
125        split: The choice of data split.
126        resize_inputs: Whether to resize the inputs to the patch shape.
127        download: Whether to download the data if it is not present.
128        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
129
130    Returns:
131        The segmentation dataset.
132    """
133    image_paths, gt_paths = get_bagls_vf_paths(path, split, download)
134
135    if resize_inputs:
136        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
137        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
138            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
139        )
140
141    return torch_em.default_segmentation_dataset(
142        raw_paths=image_paths,
143        raw_key=None,
144        label_paths=gt_paths,
145        label_key=None,
146        patch_shape=patch_shape,
147        is_seg_dataset=False,
148        **kwargs
149    )
150
151
152def get_bagls_vf_loader(
153    path: Union[os.PathLike, str],
154    patch_shape: Tuple[int, int],
155    batch_size: int,
156    split: Literal["train", "test"],
157    resize_inputs: bool = False,
158    download: bool = False,
159    **kwargs
160) -> DataLoader:
161    """Get the BAGLS-VF dataloader for vocal fold and glottal area segmentation.
162
163    Args:
164        path: Filepath to a folder where the data is downloaded for further processing.
165        patch_shape: The patch shape to use for training.
166        batch_size: The batch size for training.
167        split: The choice of data split.
168        resize_inputs: Whether to resize the inputs to the patch shape.
169        download: Whether to download the data if it is not present.
170        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
171
172    Returns:
173        The DataLoader.
174    """
175    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
176    dataset = get_bagls_vf_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
177    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)