torch_em.data.datasets.medical.orvs

  1"""The ORVS dataset contains annotations for retinal vessel segmentation in fundus images.
  2
  3It comprises 49 high-resolution fundus images (42 for training and 7 for testing) collected
  4at a clinic in Calgary, Canada, with vessel masks manually traced by a trained expert.
  5
  6NOTE: The original repository (AbdullahSarhan/ICPRVessels) referenced by the publication has
  7been removed from GitHub. This module downloads the dataset from a public fork of it instead:
  8https://github.com/hitszsyl/ICPRVessels (no explicit license is stated in the repository).
  9This dataset is from the publication https://doi.org/10.48550/arXiv.2012.09250.
 10Please cite it if you use this dataset in your research.
 11
 12NOTE: The labels are stored as JPEG images, so the (originally binary) vessel masks have lossy
 13compression artifacts near the mask boundaries. This module binarizes them with a fixed
 14intensity threshold when caching the labels to disk.
 15"""
 16
 17import os
 18from glob import glob
 19from tqdm import tqdm
 20from natsort import natsorted
 21from typing import Union, Tuple, Literal, List
 22
 23import numpy as np
 24import imageio.v3 as imageio
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL = "https://github.com/hitszsyl/ICPRVessels/archive/421886be66a400d88ef6fe9c9e047d586486eee0.zip"
 34CHECKSUM = "45f8c7fc637c6e302d2cf6c71251aeb53d0a08f4411d5e60bd7ae6753acdb24b"
 35
 36SPLIT_DIRS = {"train": "Train", "test": "Test"}
 37
 38
 39def get_orvs_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 40    """Download the ORVS dataset.
 41
 42    Args:
 43        path: Filepath to a folder where the data is downloaded for further processing.
 44        download: Whether to download the data if it is not present.
 45
 46    Returns:
 47        Filepath to the folder with the downloaded ORVS images and vessel annotations.
 48    """
 49    data_dir = os.path.join(
 50        path, "ICPRVessels-421886be66a400d88ef6fe9c9e047d586486eee0", "Vessels-Datasets", "ORVS"
 51    )
 52    if os.path.exists(data_dir):
 53        return data_dir
 54
 55    os.makedirs(path, exist_ok=True)
 56
 57    zip_path = os.path.join(path, "ICPRVessels.zip")
 58    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 59    util.unzip(zip_path=zip_path, dst=path)
 60
 61    return data_dir
 62
 63
 64def get_orvs_paths(
 65    path: Union[os.PathLike, str], split: Literal["train", "test"] = "train", download: bool = False
 66) -> Tuple[List[str], List[str]]:
 67    """Get paths to the ORVS data.
 68
 69    Args:
 70        path: Filepath to a folder where the data is downloaded for further processing.
 71        split: The choice of data split.
 72        download: Whether to download the data if it is not present.
 73
 74    Returns:
 75        List of filepaths for the image data.
 76        List of filepaths for the label data.
 77    """
 78    data_dir = get_orvs_data(path=path, download=download)
 79
 80    if split not in SPLIT_DIRS:
 81        raise ValueError(f"'{split}' is not a valid split.")
 82
 83    split_dir = os.path.join(data_dir, SPLIT_DIRS[split], "Original")
 84    image_paths = natsorted(glob(os.path.join(split_dir, "Images", "*.jpg")))
 85
 86    label_dir = os.path.join(split_dir, "Labels_binary")
 87    os.makedirs(label_dir, exist_ok=True)
 88
 89    gt_paths = []
 90    for image_path in tqdm(image_paths, desc="Preprocessing ORVS labels"):
 91        fname = os.path.splitext(os.path.basename(image_path))[0]
 92        gt_path = os.path.join(label_dir, f"{fname}.tif")
 93        gt_paths.append(gt_path)
 94        if os.path.exists(gt_path):
 95            continue
 96
 97        label_path = os.path.join(split_dir, "Labels", f"{fname}.jpg")
 98        label = imageio.imread(label_path)
 99        label = (np.asarray(label) > 127).astype("uint8")
100        if label.ndim == 3:
101            label = label.max(axis=-1)
102        imageio.imwrite(gt_path, label, compression="zlib")
103
104    return image_paths, gt_paths
105
106
107def get_orvs_dataset(
108    path: Union[os.PathLike, str],
109    patch_shape: Tuple[int, int],
110    split: Literal["train", "test"] = "train",
111    resize_inputs: bool = False,
112    download: bool = False,
113    **kwargs
114) -> Dataset:
115    """Get the ORVS dataset for retinal vessel segmentation in fundus images.
116
117    Args:
118        path: Filepath to a folder where the data is downloaded for further processing.
119        patch_shape: The patch shape to use for training.
120        split: The choice of data split.
121        resize_inputs: Whether to resize the inputs to the patch shape.
122        download: Whether to download the data if it is not present.
123        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
124
125    Returns:
126        The segmentation dataset.
127    """
128    image_paths, gt_paths = get_orvs_paths(path, split, download)
129
130    if resize_inputs:
131        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
132        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
133            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
134        )
135
136    return torch_em.default_segmentation_dataset(
137        raw_paths=image_paths,
138        raw_key=None,
139        label_paths=gt_paths,
140        label_key=None,
141        patch_shape=patch_shape,
142        is_seg_dataset=False,
143        **kwargs
144    )
145
146
147def get_orvs_loader(
148    path: Union[os.PathLike, str],
149    batch_size: int,
150    patch_shape: Tuple[int, int],
151    split: Literal["train", "test"] = "train",
152    resize_inputs: bool = False,
153    download: bool = False,
154    **kwargs
155) -> DataLoader:
156    """Get the ORVS dataloader for retinal vessel segmentation in fundus images.
157
158    Args:
159        path: Filepath to a folder where the data is downloaded for further processing.
160        batch_size: The batch size for training.
161        patch_shape: The patch shape to use for training.
162        split: The choice of data split.
163        resize_inputs: Whether to resize the inputs to the patch shape.
164        download: Whether to download the data if it is not present.
165        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
166
167    Returns:
168        The DataLoader.
169    """
170    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
171    dataset = get_orvs_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
172    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)