torch_em.data.datasets.medical.riga_plus

  1"""The RIGA+ dataset contains annotations for optic disc and optic cup segmentation in fundus images,
  2gathered from the original RIGA dataset (BinRushed, Magrabia and three MESSIDOR subsets), for the task
  3of glaucoma assessment and unsupervised domain adaptation.
  4
  5The dataset is hosted on Zenodo at https://zenodo.org/records/6325549 ("RIGA+ Dataset for Unsupervised
  6Domain Adaptation in Medical Image Segmentation").
  7
  8NOTE: There is a second, differently-scoped dataset that also carries the "RIGA+" name, hosted at
  9https://zenodo.org/records/8009107 ("A Fundus Image Dataset for Domain Generalization in Joint
 10Segmentation of Optic Disc and Optic Cup"). It combines images from REFUGE, Drishti-GS, ORIGA and RIGA
 11into a single collection and is not covered by this module.
 12
 13The five domains provided by this dataset (BinRushed, Magrabia, MESSIDOR_Base1, MESSIDOR_Base2 and
 14MESSIDOR_Base3) each ship six independent optic disc / cup annotations per image (one per rater). The
 15label masks are grayscale images with 3 pixel values: 0 (optic cup), 128 (optic disc, excluding the
 16cup) and 255 (background).
 17
 18Please cite the dataset if you use it for your research.
 19"""
 20
 21import os
 22from typing import Union, Tuple, Literal, List
 23
 24import pandas as pd
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL = "https://zenodo.org/records/6325549/files/RIGAPlus.zip"
 34CHECKSUM = "f9fa96df2faa70852ccd4c73d0e79ed6eb258b5a37e72292f5efa3f81f195d20"
 35
 36DOMAINS = ["BinRushed", "Magrabia", "MESSIDOR_Base1", "MESSIDOR_Base2", "MESSIDOR_Base3"]
 37
 38
 39def get_riga_plus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 40    """Download the RIGA+ 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    data_dir = os.path.join(path, "RIGA")
 50    if os.path.exists(data_dir):
 51        return path
 52
 53    os.makedirs(path, exist_ok=True)
 54
 55    zip_path = os.path.join(path, "RIGAPlus.zip")
 56    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 57    util.unzip(zip_path=zip_path, dst=path)
 58
 59    return path
 60
 61
 62def get_riga_plus_paths(
 63    path: Union[os.PathLike, str],
 64    split: Literal["train", "test"],
 65    domain: Union[str, List[str]] = DOMAINS,
 66    rater: int = 1,
 67    download: bool = False,
 68) -> Tuple[List[str], List[str]]:
 69    """Get paths to the RIGA+ data.
 70
 71    Args:
 72        path: Filepath to a folder where the data is downloaded for further processing.
 73        split: The choice of data split, as provided by the official 'train' / 'test' csv files.
 74        domain: The choice of domain(s) to use. One or several of 'BinRushed', 'Magrabia',
 75            'MESSIDOR_Base1', 'MESSIDOR_Base2' and 'MESSIDOR_Base3'.
 76        rater: The choice of rater (1 to 6) for the ground-truth optic disc / cup masks.
 77        download: Whether to download the data if it is not present.
 78
 79    Returns:
 80        List of filepaths for the image data.
 81        List of filepaths for the label data.
 82    """
 83    root_dir = get_riga_plus_data(path=path, download=download)
 84
 85    assert split in ["train", "test"], f"'{split}' is not a valid split."
 86    assert 1 <= rater <= 6, f"'{rater}' is not a valid rater choice, must be in [1, 6]."
 87
 88    domains = [domain] if isinstance(domain, str) else domain
 89    for d in domains:
 90        assert d in DOMAINS, f"'{d}' is not a valid domain, must be one of {DOMAINS}."
 91
 92    image_paths, gt_paths = [], []
 93    for d in domains:
 94        csv_path = os.path.join(root_dir, f"{d}_{split}.csv")
 95        df = pd.read_csv(csv_path)
 96        for image_rel_path, mask_rel_path in zip(df["image"], df["mask"]):
 97            stem, ext = os.path.splitext(mask_rel_path)
 98            image_paths.append(os.path.join(root_dir, image_rel_path))
 99            gt_paths.append(os.path.join(root_dir, f"{stem}-{rater}{ext}"))
100
101    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0
102    for image_path, gt_path in zip(image_paths, gt_paths):
103        assert os.path.exists(image_path), image_path
104        assert os.path.exists(gt_path), gt_path
105
106    return image_paths, gt_paths
107
108
109def get_riga_plus_dataset(
110    path: Union[os.PathLike, str],
111    patch_shape: Tuple[int, int],
112    split: Literal["train", "test"],
113    domain: Union[str, List[str]] = DOMAINS,
114    rater: int = 1,
115    resize_inputs: bool = False,
116    download: bool = False,
117    **kwargs
118) -> Dataset:
119    """Get the RIGA+ dataset for segmentation of optic disc and optic cup 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        split: The choice of data split.
125        domain: The choice of domain(s) to use.
126        rater: The choice of rater for the ground-truth optic disc / cup masks.
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, gt_paths = get_riga_plus_paths(path, split, domain, rater, 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=gt_paths,
146        label_key=None,
147        patch_shape=patch_shape,
148        is_seg_dataset=False,
149        **kwargs
150    )
151
152
153def get_riga_plus_loader(
154    path: Union[os.PathLike, str],
155    batch_size: int,
156    patch_shape: Tuple[int, int],
157    split: Literal["train", "test"],
158    domain: Union[str, List[str]] = DOMAINS,
159    rater: int = 1,
160    resize_inputs: bool = False,
161    download: bool = False,
162    **kwargs
163) -> DataLoader:
164    """Get the RIGA+ dataloader for segmentation of optic disc and optic cup 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        split: The choice of data split.
171        domain: The choice of domain(s) to use.
172        rater: The choice of rater for the ground-truth optic disc / cup masks.
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_riga_plus_dataset(path, patch_shape, split, domain, rater, resize_inputs, download, **ds_kwargs)
182    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)