torch_em.data.datasets.medical.hrf

  1"""The HRF dataset contains annotations for retinal vessel segmentation in high-resolution
  2fundus images of healthy, diabetic retinopathy and glaucomatous eyes.
  3
  4This dataset is located at https://www5.cs.fau.de/research/data/fundus-images/.
  5The dataset is from the publication https://doi.org/10.1155/2013/154860.
  6The dataset is licensed under CC BY 4.0 (see https://www5.cs.fau.de/research/data/fundus-images/
  7for details). Please cite the publication above if you use this dataset for your research.
  8"""
  9
 10import os
 11from glob import glob
 12from typing import Union, Tuple, Literal, List
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21BASE_URL = "https://www5.cs.fau.de/fileadmin/research/datasets/fundus-images/"
 22
 23CATEGORIES = {
 24    "healthy": "healthy",
 25    "diabetic_retinopathy": "diabetic_retinopathy",
 26    "glaucoma": "glaucoma",
 27}
 28
 29CHECKSUMS = {
 30    "healthy": "a4ce863b87371cecca8c841c4f7d8b06e39c5298fc78b103c35781bdb7eff389",
 31    "diabetic_retinopathy": "278630bea52c2096dbfa4b6486647ea60e56db1d67f1308dab2ff4383481a244",
 32    "glaucoma": "c3428a5eb971ce3165f21463c3e6bf0758531e30035303236b785e6b2d005071",
 33    "healthy_manualsegm": "2f8c67ee83e9ba16707119aa9ad98ad978c70566d467fe84b328854528ef16b8",
 34    "diabetic_retinopathy_manualsegm": "5d949ed31d4b825f33b7b14123d21f2562588c8bcbaf33984ddd51ed706d6dbd",
 35    "glaucoma_manualsegm": "ad3a0a39d2226a66da2c4072e031e23be9d27c8804e9a314db3ae757fc6f910a",
 36}
 37
 38
 39def get_hrf_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 40    """Download the HRF 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 where the data is downloaded.
 48    """
 49    os.makedirs(path, exist_ok=True)
 50
 51    for category in CATEGORIES:
 52        image_dir = os.path.join(path, category)
 53        if not os.path.exists(image_dir):
 54            zip_path = os.path.join(path, f"{category}.zip")
 55            util.download_source(
 56                path=zip_path, url=f"{BASE_URL}{category}.zip", download=download, checksum=CHECKSUMS[category],
 57            )
 58            util.unzip(zip_path=zip_path, dst=image_dir)
 59
 60        label_dir = os.path.join(path, f"{category}_manualsegm")
 61        if not os.path.exists(label_dir):
 62            zip_path = os.path.join(path, f"{category}_manualsegm.zip")
 63            util.download_source(
 64                path=zip_path,
 65                url=f"{BASE_URL}{category}_manualsegm.zip",
 66                download=download,
 67                checksum=CHECKSUMS[f"{category}_manualsegm"],
 68            )
 69            util.unzip(zip_path=zip_path, dst=label_dir)
 70
 71    return path
 72
 73
 74def get_hrf_paths(
 75    path: Union[os.PathLike, str],
 76    category: Literal["healthy", "diabetic_retinopathy", "glaucoma", "all"] = "all",
 77    download: bool = False,
 78) -> Tuple[List[str], List[str]]:
 79    """Get paths to the HRF data.
 80
 81    Args:
 82        path: Filepath to a folder where the data is downloaded for further processing.
 83        category: The choice of eye category. One of 'healthy', 'diabetic_retinopathy', 'glaucoma' or
 84            'all' (uses all three categories).
 85        download: Whether to download the data if it is not present.
 86
 87    Returns:
 88        List of filepaths for the image data.
 89        List of filepaths for the label data.
 90    """
 91    if category == "all":
 92        categories = list(CATEGORIES)
 93    elif category in CATEGORIES:
 94        categories = [category]
 95    else:
 96        raise ValueError(f"'{category}' is not a valid category. Choose from {list(CATEGORIES) + ['all']}.")
 97
 98    data_dir = get_hrf_data(path=path, download=download)
 99
100    image_paths, label_paths = [], []
101    for cat in categories:
102        cat_image_paths = sorted(
103            glob(os.path.join(data_dir, cat, "*.jpg")) + glob(os.path.join(data_dir, cat, "*.JPG"))
104        )
105        for image_path in cat_image_paths:
106            fname = os.path.splitext(os.path.basename(image_path))[0]
107            label_path = os.path.join(data_dir, f"{cat}_manualsegm", f"{fname}.tif")
108            assert os.path.exists(label_path), f"The label at '{label_path}' does not exist."
109            image_paths.append(image_path)
110            label_paths.append(label_path)
111
112    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
113
114    return image_paths, label_paths
115
116
117def get_hrf_dataset(
118    path: Union[os.PathLike, str],
119    patch_shape: Tuple[int, int],
120    category: Literal["healthy", "diabetic_retinopathy", "glaucoma", "all"] = "all",
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> Dataset:
125    """Get the HRF dataset for segmentation of retinal blood vessels in high-resolution fundus images.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        patch_shape: The patch shape to use for training.
130        category: The choice of eye category.
131        resize_inputs: Whether to resize the inputs to the expected patch shape.
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
134
135    Returns:
136        The segmentation dataset.
137    """
138    image_paths, label_paths = get_hrf_paths(path=path, category=category, download=download)
139
140    if resize_inputs:
141        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
142        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
143            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
144        )
145
146    return torch_em.default_segmentation_dataset(
147        raw_paths=image_paths,
148        raw_key=None,
149        label_paths=label_paths,
150        label_key=None,
151        patch_shape=patch_shape,
152        is_seg_dataset=False,
153        **kwargs
154    )
155
156
157def get_hrf_loader(
158    path: Union[os.PathLike, str],
159    batch_size: int,
160    patch_shape: Tuple[int, int],
161    category: Literal["healthy", "diabetic_retinopathy", "glaucoma", "all"] = "all",
162    resize_inputs: bool = False,
163    download: bool = False,
164    **kwargs
165) -> DataLoader:
166    """Get the HRF dataloader for segmentation of retinal blood vessels in high-resolution fundus images.
167
168    Args:
169        path: Filepath to a folder where the data is downloaded for further processing.
170        batch_size: The batch size for training.
171        patch_shape: The patch shape to use for training.
172        category: The choice of eye category.
173        resize_inputs: Whether to resize the inputs to the expected patch shape.
174        download: Whether to download the data if it is not present.
175        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
176
177    Returns:
178        The DataLoader.
179    """
180    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
181    dataset = get_hrf_dataset(path, patch_shape, category, resize_inputs, download, **ds_kwargs)
182    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)