torch_em.data.datasets.medical.les_av

  1"""The LES-AV dataset contains annotations for artery-vein segmentation in fundus images,
  2paired with glaucoma diagnosis information.
  3
  4For the class labels: red represents arteries, blue represents veins, green represents artery-vein
  5crossings, and white represents vessels of uncertain classification.
  6
  7This dataset is located at https://figshare.com/articles/dataset/LES-AV_dataset/11857698.
  8The dataset is from the publication https://doi.org/10.1007/978-3-030-00934-2_8.
  9Please cite it if you use this dataset for your research.
 10"""
 11
 12import os
 13from glob import glob
 14from pathlib import Path
 15from typing import Union, Tuple, List
 16
 17import numpy as np
 18import imageio.v3 as imageio
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://ndownloader.figshare.com/files/21732282"
 28CHECKSUM = "6ee8186cfe68350e212dcd2de5c59d7f23a21988e9e01d796eee166282c31350"
 29
 30
 31def _process_labels(data_dir):
 32    label_paths = glob(os.path.join(data_dir, "arteries-and-veins", "*.png"))
 33    for label_path in label_paths:
 34        labels = imageio.imread(label_path)
 35
 36        # New empty label.
 37        neu_labels = np.zeros(labels.shape[:2])
 38
 39        # Map labels to specific ids.
 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_les_av_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 51    """Download the LES-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, "LES-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, "LES-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_les_av_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 76    """Get paths to the LES-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_les_av_data(path=path, download=download)
 87
 88    image_paths = sorted(glob(os.path.join(data_dir, "images", "*.png")))
 89    label_paths = sorted(glob(os.path.join(data_dir, "arteries-and-veins", "*.tif")))
 90
 91    return image_paths, label_paths
 92
 93
 94def get_les_av_dataset(
 95    path: Union[os.PathLike, str],
 96    patch_shape: Tuple[int, int],
 97    resize_inputs: bool = False,
 98    download: bool = False,
 99    **kwargs
100) -> Dataset:
101    """Get the LES-AV dataset for artery-vein segmentation in fundus images.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        patch_shape: The patch shape to use for training.
106        resize_inputs: Whether to resize the inputs to the patch shape.
107        download: Whether to download the data if it is not present.
108        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
109
110    Returns:
111        The segmentation dataset.
112    """
113    image_paths, label_paths = get_les_av_paths(path=path, download=download)
114
115    if resize_inputs:
116        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
117        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
118            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
119        )
120
121    return torch_em.default_segmentation_dataset(
122        raw_paths=image_paths,
123        raw_key=None,
124        label_paths=label_paths,
125        label_key=None,
126        patch_shape=patch_shape,
127        is_seg_dataset=False,
128        **kwargs
129    )
130
131
132def get_les_av_loader(
133    path: Union[os.PathLike, str],
134    batch_size: int,
135    patch_shape: Tuple[int, int],
136    resize_inputs: bool = False,
137    download: bool = False,
138    **kwargs
139) -> DataLoader:
140    """Get the LES-AV dataloader for artery-vein segmentation in fundus images.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        batch_size: The batch size for training.
145        patch_shape: The patch shape to use for training.
146        resize_inputs: Whether to resize the inputs to the patch shape.
147        download: Whether to download the data if it is not present.
148        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
149
150    Returns:
151        The DataLoader.
152    """
153    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
154    dataset = get_les_av_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
155    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)