torch_em.data.datasets.medical.papila

  1"""The Papila dataset contains annotations for optic disc and optic cup
  2segmentation in Fundus images.
  3
  4This dataset is located at https://figshare.com/articles/dataset/PAPILA/14798004/2.
  5The dataset is from the publication https://doi.org/10.1038/s41597-022-01388-1.
  6Please cite it if you use this dataset for your research.
  7"""
  8
  9import os
 10from glob import glob
 11from tqdm import tqdm
 12from pathlib import Path
 13from typing import Union, Tuple, Literal, List
 14
 15import numpy as np
 16from skimage import draw
 17import imageio.v3 as imageio
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URL = "https://figshare.com/ndownloader/files/35013982"
 25CHECKSUM = "15b053dff496bc8e53eb8a8d0707ef73ba3d56c988eea92b65832c9c82852a7d"
 26
 27
 28def get_papila_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 29    """Download the Papila dataset.
 30
 31    Args:
 32        path: Filepath to a folder where the data is downloaded for further processing.
 33        download: Whether to download the data if it is not present.
 34
 35    Returns:
 36        Filepath where the data is downloaded.
 37    """
 38    data_dir = os.path.join(path, "PapilaDB-PAPILA-17f8fa7746adb20275b5b6a0d99dc9dfe3007e9f")
 39    if os.path.exists(data_dir):
 40        return data_dir
 41
 42    os.makedirs(path, exist_ok=True)
 43
 44    zip_path = os.path.join(path, "papila.zip")
 45    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 46    util.unzip(zip_path=zip_path, dst=path)
 47
 48    return data_dir
 49
 50
 51# contour_to_mask() functions taken from https://github.com/matterport/Mask_RCNN
 52def contour_to_mask(cont, img_shape):
 53    """Return mask given a contour and the shape of image
 54    """
 55    c = np.loadtxt(cont)
 56    mask = np.zeros(img_shape[:-1], dtype=np.uint8)
 57    rr, cc = draw.polygon(c[:, 1], c[:, 0])
 58    mask[rr, cc] = 1
 59    return mask
 60
 61
 62def _preprocess_labels(data_dir, image_paths, task, expert_choice):
 63    gt_dir = os.path.join(data_dir, "ground_truth")
 64    os.makedirs(gt_dir, exist_ok=True)
 65
 66    patient_ids = [Path(image_path).stem for image_path in image_paths]
 67
 68    input_shape = (1934, 2576, 3)  # shape of the input images
 69    gt_paths = []
 70    for patient_id in tqdm(patient_ids, desc=f"Converting contours to segmentations for '{expert_choice}'"):
 71        gt_contours = sorted(
 72            glob(os.path.join(data_dir, "ExpertsSegmentations", "Contours", f"{patient_id}_{task}_{expert_choice}.txt"))
 73        )
 74
 75        for gt_contour in gt_contours:
 76            tmp_task = Path(gt_contour).stem.split("_")[1]
 77            gt_path = os.path.join(gt_dir, f"{patient_id}_{tmp_task}_{expert_choice}.tif")
 78            gt_paths.append(gt_path)
 79            if os.path.exists(gt_path):
 80                continue
 81
 82            semantic_labels = contour_to_mask(cont=gt_contour, img_shape=input_shape)
 83            imageio.imwrite(gt_path, semantic_labels)
 84
 85    return gt_paths
 86
 87
 88def get_papila_paths(
 89    path: Union[os.PathLike, str],
 90    split: Literal['train', 'val', 'test'],
 91    task: Literal["cup", "disc"] = "disc",
 92    expert_choice: Literal["exp1", "exp2"] = "exp1",
 93    download: bool = False
 94) -> Tuple[List[str], List[str]]:
 95    """Get paths to the Papila dataset.
 96
 97    Args:
 98        path: Filepath to a folder where the data is downloaded for further processing.
 99        split: The choice of data split.
100        task: The choice of labels for specific task.
101        expert_choice: The choice of expert annotator.
102        download: Whether to download the data if it is not present.
103
104    Returns:
105        List of filepaths for the image data.
106        List of filepaths for the label data.
107    """
108    data_dir = get_papila_data(path=path, download=download)
109
110    assert expert_choice in ["exp1", "exp2"], f"'{expert_choice}' is not a valid expert choice."
111    assert task in ["cup", "disc"], f"'{task}' is not a valid task."
112
113    image_paths = sorted(glob(os.path.join(data_dir, "FundusImages", "*.jpg")))
114    gt_paths = _preprocess_labels(data_dir, image_paths, task, expert_choice)
115
116    if split == "train":
117        image_paths, gt_paths = image_paths[:350], gt_paths[:350]
118    elif split == "val":
119        image_paths, gt_paths = image_paths[350:400], gt_paths[350:400]
120    elif split == "test":
121        image_paths, gt_paths = image_paths[400:], gt_paths[400:]
122    else:
123        raise ValueError(f"'{split}' is not a valid split.")
124
125    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0
126
127    return image_paths, gt_paths
128
129
130def get_papila_dataset(
131    path: Union[os.PathLike, str],
132    patch_shape: Tuple[int, int],
133    split: Literal['train', 'val', 'test'],
134    task: Literal["cup", "disc"] = "disc",
135    expert_choice: Literal["exp1", "exp2"] = "exp1",
136    resize_inputs: bool = False,
137    download: bool = False,
138    **kwargs
139):
140    """Get the Papila dataset for segmentation of optic cup and optic disc in fundus images.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        patch_shape: The patch shape to use for training.
145        split: Te choice of data split.
146        task: The choice of labels for specific task.
147        expert_choice: The choice of expert annotator.
148        resize_inputs: Whether to resize the inputs to the expected patch shape.
149        download: Whether to download the data if it is not present.
150        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
151
152    Returns:
153        The segmentation dataset.
154    """
155    image_paths, gt_paths = get_papila_paths(path, split, task, expert_choice, download)
156
157    if resize_inputs:
158        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
159        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
160            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
161        )
162
163    dataset = torch_em.default_segmentation_dataset(
164        raw_paths=image_paths,
165        raw_key=None,
166        label_paths=gt_paths,
167        label_key=None,
168        patch_shape=patch_shape,
169        is_seg_dataset=False,
170        **kwargs
171    )
172
173    return dataset
174
175
176def get_papila_loader(
177    path: Union[os.PathLike, str],
178    batch_size: int,
179    patch_shape: Tuple[int, int],
180    split: Literal['train', 'val', 'test'],
181    task: Literal["cup", "disc"] = "disc",
182    expert_choice: Literal["exp1", "exp2"] = "exp1",
183    resize_inputs: bool = False,
184    download: bool = False,
185    **kwargs
186):
187    """Get the Papila dataloader for segmentation of optic cup and optic disc in fundus images.
188
189    Args:
190        path: Filepath to a folder where the data is downloaded for further processing.
191        batch_size: The batch size for training.
192        patch_shape: The patch shape to use for training.
193        split: The choice of data split.
194        task: The choice of labels for specific task.
195        expert_choice: The choice of expert annotator.
196        resize_inputs: Whether to resize the inputs to the expected patch shape.
197        download: Whether to download the data if it is not present.
198        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
199
200    Returns:
201        The DataLoader.
202    """
203    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
204    dataset = get_papila_dataset(path, patch_shape, split, task, expert_choice, resize_inputs, download, **ds_kwargs)
205    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)