torch_em.data.datasets.medical.rossa

  1"""The ROSSA dataset contains annotations for retinal vessel segmentation in OCTA images.
  2
  3It comprises 918 OCTA images: 300 images with manually verified vessel masks (split into
  4'train_manual', 'val' and 'test' folders of 100 images each) and 618 further images with
  5vessel masks that were generated semi-automatically with the Segment Anything Model, stored
  6in the 'train_sam' folder. The manually annotated subset is the reliable one for evaluation
  7and is used by default; the SAM-assisted subset trades annotation quality for scale.
  8
  9The dataset is located at https://github.com/nhjydywd/OCTA-FRNet (MIT license).
 10This dataset is from the publication https://doi.org/10.48550/arXiv.2309.09483.
 11Please cite it if you use this dataset in your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 17from typing import Union, Tuple, Literal, List
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26URL = "https://github.com/nhjydywd/OCTA-FRNet/archive/53e87f8c9b9392c1dcababea07e986f111f4017c.zip"
 27CHECKSUM = "2d408a9b124aadbb5ebdf49f2aeb5e24c0fee2ef77e6a0be06648f12df8633ca"
 28
 29ANNOTATION_DIRS = {
 30    "manual": ["train_manual", "val", "test"],
 31    "sam_assisted": ["train_sam"],
 32}
 33
 34
 35def get_rossa_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 36    """Download the ROSSA dataset.
 37
 38    Args:
 39        path: Filepath to a folder where the data is downloaded for further processing.
 40        download: Whether to download the data if it is not present.
 41
 42    Returns:
 43        Filepath to the folder with the downloaded ROSSA images and vessel masks.
 44    """
 45    data_dir = os.path.join(path, "OCTA-FRNet-53e87f8c9b9392c1dcababea07e986f111f4017c", "dataset", "ROSSA")
 46    if os.path.exists(data_dir):
 47        return data_dir
 48
 49    os.makedirs(path, exist_ok=True)
 50
 51    zip_path = os.path.join(path, "OCTA-FRNet.zip")
 52    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 53    util.unzip(zip_path=zip_path, dst=path)
 54
 55    return data_dir
 56
 57
 58def get_rossa_paths(
 59    path: Union[os.PathLike, str],
 60    annotation: Literal["manual", "sam_assisted", "all"] = "manual",
 61    download: bool = False,
 62) -> Tuple[List[str], List[str]]:
 63    """Get paths to the ROSSA data.
 64
 65    Args:
 66        path: Filepath to a folder where the data is downloaded for further processing.
 67        annotation: The choice of annotation source. Either the 300 manually annotated images
 68            ('manual', the default), the 618 SAM-assisted images ('sam_assisted') or both ('all').
 69        download: Whether to download the data if it is not present.
 70
 71    Returns:
 72        List of filepaths for the image data.
 73        List of filepaths for the label data.
 74    """
 75    data_dir = get_rossa_data(path=path, download=download)
 76
 77    if annotation == "all":
 78        sub_dirs = ANNOTATION_DIRS["manual"] + ANNOTATION_DIRS["sam_assisted"]
 79    elif annotation in ANNOTATION_DIRS:
 80        sub_dirs = ANNOTATION_DIRS[annotation]
 81    else:
 82        raise ValueError(f"'{annotation}' is not a valid annotation choice.")
 83
 84    image_paths, gt_paths = [], []
 85    for sub_dir in sub_dirs:
 86        image_paths.extend(natsorted(glob(os.path.join(data_dir, sub_dir, "image", "*.png"))))
 87        gt_paths.extend(natsorted(glob(os.path.join(data_dir, sub_dir, "label", "*.png"))))
 88
 89    return image_paths, gt_paths
 90
 91
 92def get_rossa_dataset(
 93    path: Union[os.PathLike, str],
 94    patch_shape: Tuple[int, int],
 95    annotation: Literal["manual", "sam_assisted", "all"] = "manual",
 96    resize_inputs: bool = False,
 97    download: bool = False,
 98    **kwargs
 99) -> Dataset:
100    """Get the ROSSA dataset for retinal vessel segmentation in OCTA images.
101
102    Args:
103        path: Filepath to a folder where the data is downloaded for further processing.
104        patch_shape: The patch shape to use for training.
105        annotation: The choice of annotation source. Either the 300 manually annotated images
106            ('manual', the default), the 618 SAM-assisted images ('sam_assisted') or both ('all').
107        resize_inputs: Whether to resize the inputs.
108        download: Whether to download the data if it is not present.
109        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
110
111    Returns:
112        The segmentation dataset.
113    """
114    image_paths, gt_paths = get_rossa_paths(path, annotation, download)
115
116    if resize_inputs:
117        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
118        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
119            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
120        )
121
122    return torch_em.default_segmentation_dataset(
123        raw_paths=image_paths,
124        raw_key=None,
125        label_paths=gt_paths,
126        label_key=None,
127        patch_shape=patch_shape,
128        is_seg_dataset=False,
129        **kwargs
130    )
131
132
133def get_rossa_loader(
134    path: Union[os.PathLike, str],
135    batch_size: int,
136    patch_shape: Tuple[int, int],
137    annotation: Literal["manual", "sam_assisted", "all"] = "manual",
138    resize_inputs: bool = False,
139    download: bool = False,
140    **kwargs
141) -> DataLoader:
142    """Get the ROSSA dataloader for retinal vessel segmentation in OCTA images.
143
144    Args:
145        path: Filepath to a folder where the data is downloaded for further processing.
146        batch_size: The batch size for training.
147        patch_shape: The patch shape to use for training.
148        annotation: The choice of annotation source. Either the 300 manually annotated images
149            ('manual', the default), the 618 SAM-assisted images ('sam_assisted') or both ('all').
150        resize_inputs: Whether to resize the inputs.
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
153
154    Returns:
155        The DataLoader.
156    """
157    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
158    dataset = get_rossa_dataset(path, patch_shape, annotation, resize_inputs, download, **ds_kwargs)
159    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)