torch_em.data.datasets.medical.puwf_av

  1"""The PUWF-AV dataset contains annotations for artery-vein segmentation in pediatric
  2ultra-widefield (UWF) fundus images.
  3
  4For the class labels: red represents arteries, blue represents veins, green represents artery-vein
  5crossings, and white represents vessels of uncertain classification.
  6
  7NOTE: The dataset does not ship an official train / val / test split, so `get_puwf_av_paths` returns
  8all 60 images.
  9
 10This dataset is from the publication https://doi.org/10.1038/s41597-026-08342-5.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from pathlib import Path
 17from natsort import natsorted
 18from typing import Union, Tuple, List
 19
 20import numpy as np
 21import imageio.v3 as imageio
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://ndownloader.figshare.com/files/65105883"
 31CHECKSUM = "d8b9a060dde125e767f999de7cbfc8d49cfac584e22611d915c40e9ceb21a054"
 32
 33
 34def _process_labels(data_dir):
 35    label_paths = glob(os.path.join(data_dir, "annotation", "*.png"))
 36    for label_path in label_paths:
 37        labels = imageio.imread(label_path)
 38
 39        neu_labels = np.zeros(labels.shape[:2])
 40        neu_labels[np.all(labels == (255, 0, 0), axis=-1)] = 1  # red are arteries.
 41        neu_labels[np.all(labels == (0, 0, 255), axis=-1)] = 2  # blue are veins.
 42        neu_labels[np.all(labels == (0, 255, 0), axis=-1)] = 3  # green are overlaps.
 43        neu_labels[np.all(labels == (255, 255, 255), axis=-1)] = 4  # white are unknown.
 44
 45        imageio.imwrite(Path(label_path).with_suffix(".tif"), neu_labels, compression="zlib")
 46
 47        os.remove(label_path)
 48
 49
 50def get_puwf_av_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 51    """Download the PUWF-AV dataset.
 52
 53    Args:
 54        path: Filepath to a folder where the data is downloaded for further processing.
 55        download: Whether to download the data if it is not present.
 56
 57    Returns:
 58        Filepath where the data is downloaded.
 59    """
 60    data_dir = os.path.join(path, "PUWF-AV")
 61    if os.path.exists(data_dir):
 62        return data_dir
 63
 64    os.makedirs(path, exist_ok=True)
 65
 66    zip_path = os.path.join(path, "PUWF-AV.zip")
 67    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 68    util.unzip(zip_path=zip_path, dst=path)
 69
 70    _process_labels(data_dir)
 71
 72    return data_dir
 73
 74
 75def get_puwf_av_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 76    """Get paths to the PUWF-AV data.
 77
 78    Args:
 79        path: Filepath to a folder where the data is downloaded for further processing.
 80        download: Whether to download the data if it is not present.
 81
 82    Returns:
 83        List of filepaths for the image data.
 84        List of filepaths for the label data.
 85    """
 86    data_dir = get_puwf_av_data(path, download)
 87
 88    raw_paths = natsorted(glob(os.path.join(data_dir, "images", "*.png")))
 89    label_paths = natsorted(glob(os.path.join(data_dir, "annotation", "*.tif")))
 90
 91    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 92
 93    return raw_paths, label_paths
 94
 95
 96def get_puwf_av_dataset(
 97    path: Union[os.PathLike, str],
 98    patch_shape: Tuple[int, int],
 99    resize_inputs: bool = False,
100    download: bool = False,
101    **kwargs
102) -> Dataset:
103    """Get the PUWF-AV dataset for artery-vein segmentation.
104
105    Args:
106        path: Filepath to a folder where the data is downloaded for further processing.
107        patch_shape: The patch shape to use for training.
108        resize_inputs: Whether to resize the inputs to the patch shape.
109        download: Whether to download the data if it is not present.
110        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
111
112    Returns:
113        The segmentation dataset.
114    """
115    raw_paths, label_paths = get_puwf_av_paths(path, download)
116
117    if resize_inputs:
118        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
119        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
120            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
121        )
122
123    return torch_em.default_segmentation_dataset(
124        raw_paths=raw_paths,
125        raw_key=None,
126        label_paths=label_paths,
127        label_key=None,
128        is_seg_dataset=False,
129        patch_shape=patch_shape,
130        ndim=2,
131        **kwargs
132    )
133
134
135def get_puwf_av_loader(
136    path: Union[os.PathLike, str],
137    batch_size: int,
138    patch_shape: Tuple[int, int],
139    resize_inputs: bool = False,
140    download: bool = False,
141    **kwargs
142) -> DataLoader:
143    """Get the PUWF-AV dataloader for artery-vein segmentation.
144
145    Args:
146        path: Filepath to a folder where the data is downloaded for further processing.
147        batch_size: The batch size for training.
148        patch_shape: The patch shape to use for training.
149        resize_inputs: Whether to resize the inputs to the patch shape.
150        download: Whether to download the data if it is not present.
151        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
152
153    Returns:
154        The DataLoader.
155    """
156    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
157    dataset = get_puwf_av_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
158    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)