torch_em.data.datasets.medical.fovea

  1"""FOVEA is a dataset for optic disc and retinal vessel segmentation in preoperative and
  2intraoperative retinal fundus images, comprising 40 patients collected at Moorfields Eye
  3Hospital (London, UK). For each patient and each domain (preoperative color fundus photography
  4and intraoperative retinal microscopy), the green channel image used for annotation as well as
  5binary optic disc and retinal vessel masks from two independent clinical research fellows are
  6provided.
  7
  8The dataset is located at https://doi.org/10.6084/m9.figshare.28329338.
  9This dataset is from the publication https://doi.org/10.1038/s41597-025-04965-2.
 10Please cite it if you use this dataset in your research.
 11"""
 12
 13import os
 14from glob import glob
 15from natsort import natsorted
 16from typing import Union, Tuple, Literal, List
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://ndownloader.figshare.com/files/52512728"
 26CHECKSUM = "4f9b88468e79b89d9c561932c28e6bd0974bde1ce2a21b1027b7dca2e09bf24d"
 27
 28DOMAINS = {"preoperative": "p", "intraoperative": "i"}
 29ANNOTATIONS = {"optic_disc": "od", "vessels": "ve"}
 30
 31
 32def get_fovea_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the FOVEA dataset.
 34
 35    Args:
 36        path: Filepath to a folder where the data is downloaded for further processing.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath where the data is downloaded.
 41    """
 42    if os.path.exists(path) and glob(os.path.join(path, "FOVEA*_img.png")):
 43        return path
 44
 45    os.makedirs(path, exist_ok=True)
 46
 47    zip_path = os.path.join(path, "FOVEA.zip")
 48    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 49    util.unzip(zip_path=zip_path, dst=path)
 50
 51    return path
 52
 53
 54def get_fovea_paths(
 55    path: Union[os.PathLike, str],
 56    domain: Literal["preoperative", "intraoperative", "both"] = "both",
 57    annotation: Literal["optic_disc", "vessels"] = "vessels",
 58    annotator: Literal[1, 2, "both"] = 1,
 59    download: bool = False,
 60) -> Tuple[List[str], List[str]]:
 61    """Get paths to the FOVEA data.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        domain: The choice of imaging domain. One of 'preoperative', 'intraoperative' or 'both'.
 66        annotation: The choice of segmentation target. Either 'optic_disc' or 'vessels'.
 67        annotator: The choice of annotator. Either 1, 2 or 'both' (uses annotations from both annotators).
 68        download: Whether to download the data if it is not present.
 69
 70    Returns:
 71        List of filepaths for the image data.
 72        List of filepaths for the label data.
 73    """
 74    data_dir = get_fovea_data(path=path, download=download)
 75
 76    if domain == "both":
 77        domain_tags = list(DOMAINS.values())
 78    elif domain in DOMAINS:
 79        domain_tags = [DOMAINS[domain]]
 80    else:
 81        raise ValueError(f"'{domain}' is not a valid domain. Choose from {list(DOMAINS) + ['both']}.")
 82
 83    if annotation not in ANNOTATIONS:
 84        raise ValueError(f"'{annotation}' is not a valid annotation. Choose from {list(ANNOTATIONS)}.")
 85    annotation_tag = ANNOTATIONS[annotation]
 86
 87    if annotator == "both":
 88        annotators = [1, 2]
 89    elif annotator in (1, 2):
 90        annotators = [annotator]
 91    else:
 92        raise ValueError(f"'{annotator}' is not a valid annotator. Choose from 1, 2 or 'both'.")
 93
 94    image_paths, label_paths = [], []
 95    for domain_tag in domain_tags:
 96        for ann in annotators:
 97            label_glob = natsorted(glob(os.path.join(data_dir, f"FOVEA*_{domain_tag}_{annotation_tag}_{ann}.png")))
 98            for label_path in label_glob:
 99                image_path = label_path.replace(f"_{annotation_tag}_{ann}.png", "_img.png")
100                assert os.path.exists(image_path), f"The image at '{image_path}' does not exist."
101                image_paths.append(image_path)
102                label_paths.append(label_path)
103
104    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
105
106    return image_paths, label_paths
107
108
109def get_fovea_dataset(
110    path: Union[os.PathLike, str],
111    patch_shape: Tuple[int, int],
112    domain: Literal["preoperative", "intraoperative", "both"] = "both",
113    annotation: Literal["optic_disc", "vessels"] = "vessels",
114    annotator: Literal[1, 2, "both"] = 1,
115    resize_inputs: bool = False,
116    download: bool = False,
117    **kwargs
118) -> Dataset:
119    """Get the FOVEA dataset for optic disc and retinal vessel segmentation in fundus images.
120
121    Args:
122        path: Filepath to a folder where the data is downloaded for further processing.
123        patch_shape: The patch shape to use for training.
124        domain: The choice of imaging domain.
125        annotation: The choice of segmentation target.
126        annotator: The choice of annotator.
127        resize_inputs: Whether to resize the inputs to the expected patch shape.
128        download: Whether to download the data if it is not present.
129        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
130
131    Returns:
132        The segmentation dataset.
133    """
134    image_paths, label_paths = get_fovea_paths(path, domain, annotation, annotator, download)
135
136    if resize_inputs:
137        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
138        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
139            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
140        )
141
142    return torch_em.default_segmentation_dataset(
143        raw_paths=image_paths,
144        raw_key=None,
145        label_paths=label_paths,
146        label_key=None,
147        patch_shape=patch_shape,
148        is_seg_dataset=False,
149        **kwargs
150    )
151
152
153def get_fovea_loader(
154    path: Union[os.PathLike, str],
155    batch_size: int,
156    patch_shape: Tuple[int, int],
157    domain: Literal["preoperative", "intraoperative", "both"] = "both",
158    annotation: Literal["optic_disc", "vessels"] = "vessels",
159    annotator: Literal[1, 2, "both"] = 1,
160    resize_inputs: bool = False,
161    download: bool = False,
162    **kwargs
163) -> DataLoader:
164    """Get the FOVEA dataloader for optic disc and retinal vessel segmentation in fundus images.
165
166    Args:
167        path: Filepath to a folder where the data is downloaded for further processing.
168        batch_size: The batch size for training.
169        patch_shape: The patch shape to use for training.
170        domain: The choice of imaging domain.
171        annotation: The choice of segmentation target.
172        annotator: The choice of annotator.
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_fovea_dataset(path, patch_shape, domain, annotation, annotator, resize_inputs, download, **ds_kwargs)
182    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)