torch_em.data.datasets.medical.idrid

  1"""The IDRID dataset contains annotations for retinal lesions and optic disc segmentation
  2in Fundus images.
  3
  4The database is located at https://ieee-dataport.org/open-access/indian-diabetic-retinopathy-image-dataset-idrid
  5The dataloader makes use of an open-source version of the original dataset hosted on Kaggle.
  6
  7The dataset is from the IDRiD challenge:
  8- https://idrid.grand-challenge.org/
  9- Porwal et al. - https://doi.org/10.1016/j.media.2019.101561
 10Please cite them if you use this dataset for your research.
 11
 12The 'refined' version (selected with `version="refined"`) uses the same 81 IDRiD images, but replaces
 13the lesion annotations with the 'Refined IDRiD' release: expert-corrected and validated annotations for
 14the four original lesion types (microaneurysms, haemorrhages, hard exudates, soft / cotton-wool exudates),
 15plus three additional proliferative DR lesion types (neovascularization, vitreous haemorrhage, intraretinal
 16microvascular abnormalities) and anatomical context (optic disc, fovea, blood vessels, retinal region), all
 17merged into a single unified multi-class label mask per image. It is located at
 18https://doi.org/10.5281/zenodo.18676805 (CC BY 4.0). The `task` argument does not apply to this version, as
 19the label masks are not split by lesion type. Please cite the dataset if you use this version:
 20https://www.mdpi.com/2306-5729/11/2/30.
 21"""
 22
 23import os
 24from glob import glob
 25from pathlib import Path
 26from typing import Union, Tuple, Literal, List
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35TASKS = {
 36    "microaneurysms": r"1. Microaneurysms",
 37    "haemorrhages": r"2. Haemorrhages",
 38    "hard_exudates": r"3. Hard Exudates",
 39    "soft_exudates": r"4. Soft Exudates",
 40    "optic_disc": r"5. Optic Disc"
 41}
 42
 43VERSIONS = ["v1", "refined"]
 44
 45URL_REFINED = {
 46    "train": "https://zenodo.org/records/18676805/files/Train.tar",
 47    "test": "https://zenodo.org/records/18676805/files/Test.tar",
 48}
 49CHECKSUM_REFINED = {
 50    "train": "cb368cbfcdcbb2a9d22b95a99301aa4a06e2d7d510a2a151d488216f570136ba",
 51    "test": "3482c23e9c7a179960ea2a831ffe665bcbc825d497403ba5d0f386642eb4d821",
 52}
 53
 54
 55def get_idrid_data(
 56    path: Union[os.PathLike, str], download: bool = False, version: Literal["v1", "refined"] = "v1"
 57) -> str:
 58    """Download the IDRID dataset.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        download: Whether to download the data if it is not present.
 63        version: The version of the dataset. Either 'v1' (the original IDRiD lesion annotations) or
 64            'refined' (the 'Refined IDRiD' release with expert-corrected and additional lesion annotations).
 65
 66    Returns:
 67        Filepath where the data is downloaded.
 68    """
 69    if version not in VERSIONS:
 70        raise ValueError(f"'{version}' is not a valid version. Please choose one of {VERSIONS}.")
 71
 72    if version == "refined":
 73        data_dir = os.path.join(path, "refined")
 74        if os.path.exists(data_dir):
 75            return data_dir
 76
 77        os.makedirs(data_dir, exist_ok=True)
 78
 79        for split, url in URL_REFINED.items():
 80            tar_path = os.path.join(path, f"{split.capitalize()}.tar")
 81            util.download_source(path=tar_path, url=url, download=download, checksum=CHECKSUM_REFINED[split])
 82            util.unzip_tarfile(tar_path=tar_path, dst=data_dir, remove=False)
 83
 84        return data_dir
 85
 86    data_dir = os.path.join(path, "data", "A.%20Segmentation")
 87    if os.path.exists(data_dir):
 88        return data_dir
 89
 90    os.makedirs(path, exist_ok=True)
 91
 92    util.download_source_kaggle(
 93        path=path, dataset_name="aaryapatel98/indian-diabetic-retinopathy-image-dataset", download=download,
 94    )
 95    zip_path = os.path.join(path, "indian-diabetic-retinopathy-image-dataset.zip")
 96    util.unzip(zip_path=zip_path, dst=os.path.join(path, "data"))
 97
 98    return data_dir
 99
100
101def get_idrid_paths(
102    path: Union[os.PathLike, str],
103    split: Literal['train', 'test'],
104    task: Literal['microaneurysms', 'haemorrhages', 'hard_exudates', 'soft_exudates', 'optic_disc'],
105    download: bool = False,
106    version: Literal["v1", "refined"] = "v1",
107) -> Tuple[List[str], List[str]]:
108    """Get paths to the IDRID data.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        split: The choice of data split.
113        task: The choice of labels for the specific task. Ignored when `version` is 'refined'.
114        download: Whether to download the data if it is not present.
115        version: The version of the dataset. Either 'v1' (the original IDRiD lesion annotations) or
116            'refined' (the 'Refined IDRiD' release with expert-corrected and additional lesion annotations).
117
118    Returns:
119        List of filepaths for the image data.
120        List of filepaths for the label data.
121    """
122    assert split in ["train", "test"]
123
124    data_dir = get_idrid_data(path=path, download=download, version=version)
125
126    if version == "refined":
127        split_dir = "Train" if split == "train" else "Test"
128        image_paths = sorted(glob(os.path.join(data_dir, split_dir, "Images", "*")))
129
130        label_dir = os.path.join(data_dir, split_dir, "Labels")
131        gt_paths = []
132        for image_path in image_paths:
133            stem = Path(image_path).stem
134            label_path = os.path.join(label_dir, f"{stem}.png")
135            if not os.path.exists(label_path):
136                label_path = os.path.join(label_dir, f"{stem}_vessel.png")
137            if not os.path.exists(label_path):
138                raise RuntimeError(f"Could not find the matching label for the image at '{image_path}'.")
139            gt_paths.append(label_path)
140
141        return image_paths, gt_paths
142
143    assert task in list(TASKS.keys())
144
145    split = r"a. Training Set" if split == "train" else r"b. Testing Set"
146    gt_paths = sorted(
147        glob(
148            os.path.join(data_dir, r"A. Segmentation", r"2. All Segmentation Groundtruths", split, TASKS[task], "*.tif")
149        )
150    )
151
152    image_dir = os.path.join(data_dir, r"A. Segmentation", r"1. Original Images", split)
153    image_paths = [os.path.join(image_dir, f"{Path(p).stem[:-3]}.jpg") for p in gt_paths]
154
155    return image_paths, gt_paths
156
157
158def get_idrid_dataset(
159    path: Union[os.PathLike, str],
160    patch_shape: Tuple[int, int],
161    split: Literal['train', 'test'],
162    task: Literal['microaneurysms', 'haemorrhages', 'hard_exudates', 'soft_exudates', 'optic_disc'] = 'optic_disc',
163    resize_inputs: bool = False,
164    download: bool = False,
165    version: Literal["v1", "refined"] = "v1",
166    **kwargs
167) -> Dataset:
168    """Get the IDRID dataset for segmentation of retinal lesions and optic disc in fundus images.
169
170    Args:
171        path: Filepath to a folder where the data is downloaded for further processing.
172        patch_shape: The patch shape to use for training.
173        split: The choice of data split.
174        task: The choice of labels for the specific task. Ignored when `version` is 'refined'.
175        resize_inputs: Whether to resize the inputs to the expected patch shape.
176        download: Whether to download the data if it is not present.
177        version: The version of the dataset. Either 'v1' (the original IDRiD lesion annotations) or
178            'refined' (the 'Refined IDRiD' release with expert-corrected and additional lesion annotations).
179        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
180
181    Returns:
182        The segmentation dataset.
183    """
184    image_paths, gt_paths = get_idrid_paths(path, split, task, download, version)
185
186    if resize_inputs:
187        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
188        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
189            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
190        )
191
192    return torch_em.default_segmentation_dataset(
193        raw_paths=image_paths,
194        raw_key=None,
195        label_paths=gt_paths,
196        label_key=None,
197        patch_shape=patch_shape,
198        is_seg_dataset=False,
199        **kwargs
200    )
201
202
203def get_idrid_loader(
204    path: Union[os.PathLike, str],
205    batch_size: int,
206    patch_shape: Tuple[int, int],
207    split: Literal['train', 'test'],
208    task: Literal['microaneurysms', 'haemorrhages', 'hard_exudates', 'soft_exudates', 'optic_disc'] = 'optic_disc',
209    resize_inputs: bool = False,
210    download: bool = False,
211    version: Literal["v1", "refined"] = "v1",
212    **kwargs
213) -> DataLoader:
214    """Get the IDRID dataloader for segmentation of retinal lesions and optic disc in fundus images.
215
216    Args:
217        path: Filepath to a folder where the data is downloaded for further processing.
218        batch_size: The batch size for training.
219        patch_shape: The patch shape to use for training.
220        split: The choice of data split.
221        task: The choice of labels for the specific task. Ignored when `version` is 'refined'.
222        resize_inputs: Whether to resize the inputs to the expected patch shape.
223        download: Whether to download the data if it is not present.
224        version: The version of the dataset. Either 'v1' (the original IDRiD lesion annotations) or
225            'refined' (the 'Refined IDRiD' release with expert-corrected and additional lesion annotations).
226        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
227
228    Returns:
229        The DataLoader.
230    """
231    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
232    dataset = get_idrid_dataset(path, patch_shape, split, task, resize_inputs, download, version, **ds_kwargs)
233    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)