torch_em.data.datasets.medical.csc_fundus_segmentation

  1"""The CSC Fundus Segmentation dataset contains annotations for subretinal fluid (SRF)
  2segmentation in fundus photographs of eyes with central serous chorioretinopathy (CSC).
  3
  4This dataset is located at https://doi.org/10.17632/4k64fwnp4k.3, under the CC BY 4.0 license.
  5The dataset is from the publication https://doi.org/10.1167/tvst.11.2.11.
  6Please cite it if you use this dataset for your research.
  7
  8The dataset comprises fundus photographs from 194 eyes with CSC, with the SRF area manually
  9segmented by three ophthalmologists (two licensed ophthalmologists and one ophthalmology
 10resident), provided here as `grader1`, `grader2` and `grader3`. Each image is shipped as a
 11paired image, with the raw fundus photograph on the left half and the corresponding
 12segmentation rendered as a pink overlay on a white background on the right half. We split
 13these paired images here into the raw image and a binarized foreground / background mask.
 14"""
 15
 16import os
 17from glob import glob
 18from pathlib import Path
 19from natsort import natsorted
 20from typing import Union, Tuple, Literal, List
 21
 22import numpy as np
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URL = "https://data.mendeley.com/public-files/datasets/4k64fwnp4k/files/{}/file_downloaded"
 33FNAMES_TO_IDS = {
 34    "Grader_1.zip": "e53c7210-9e09-4e0e-8f36-87103128fe38",
 35    "Grader_2.zip": "2b551de5-1855-47cf-88cf-6f55eb26d001",
 36    "Grader_3.zip": "fd6214bf-a33a-413c-948c-7430beebe286",
 37}
 38CHECKSUMS = {
 39    "Grader_1.zip": "e9820a9cd2b13490f972e2c58b316b0ab62dccef00564d1c1cc768e55658b435",
 40    "Grader_2.zip": "78ca3f61b58cf85ef2619b031c6c4a4a95f9af7f8bcbc567e9d295f6ead1120e",
 41    "Grader_3.zip": "a2c9699125c056cb8f46d073a72f93eec1e3a57b77370eaa15ed1a3475dba07a",
 42}
 43
 44
 45def get_csc_fundus_segmentation_data(
 46    path: Union[os.PathLike, str], grader: Literal["grader1", "grader2", "grader3"], download: bool = False
 47) -> str:
 48    """Download the CSC Fundus Segmentation dataset.
 49
 50    Args:
 51        path: Filepath to a folder where the data is downloaded for further processing.
 52        grader: The choice of annotator whose segmentations are used.
 53        download: Whether to download the data if it is not present.
 54
 55    Returns:
 56        Filepath where the data is downloaded.
 57    """
 58    assert grader in ["grader1", "grader2", "grader3"], f"'{grader}' is not a valid grader choice."
 59
 60    fname = f"Grader_{grader[-1]}.zip"
 61    data_dir = os.path.join(path, Path(fname).stem)
 62    if os.path.exists(data_dir):
 63        return data_dir
 64
 65    os.makedirs(path, exist_ok=True)
 66
 67    zip_path = os.path.join(path, fname)
 68    util.download_source(
 69        path=zip_path, url=URL.format(FNAMES_TO_IDS[fname]), download=download, checksum=CHECKSUMS[fname]
 70    )
 71    util.unzip(zip_path=zip_path, dst=data_dir)
 72
 73    return data_dir
 74
 75
 76def _split_paired_image(paired_path, image_dir, gt_dir):
 77    fname = Path(paired_path).stem
 78    image_path = os.path.join(image_dir, f"{fname}.png")
 79    gt_path = os.path.join(gt_dir, f"{fname}.png")
 80
 81    if os.path.exists(image_path) and os.path.exists(gt_path):
 82        return image_path, gt_path
 83
 84    os.makedirs(image_dir, exist_ok=True)
 85    os.makedirs(gt_dir, exist_ok=True)
 86
 87    paired_image = imageio.imread(paired_path)
 88    width = paired_image.shape[1]
 89    half = width // 2
 90
 91    raw = paired_image[:, :half]
 92
 93    # The right half renders the SRF segmentation as a pink overlay on a white background.
 94    # The seam between the concatenated halves introduces a single column of black JPEG
 95    # compression bleed at the left edge of the mask half, which we remove before thresholding.
 96    mask_rgb = paired_image[:, half:].copy()
 97    mask_rgb[:, 0] = 255
 98    mask = (np.abs(255 - mask_rgb.astype(int)).sum(axis=-1) > 30).astype("uint8")
 99
100    imageio.imwrite(image_path, raw)
101    imageio.imwrite(gt_path, mask)
102
103    return image_path, gt_path
104
105
106def get_csc_fundus_segmentation_paths(
107    path: Union[os.PathLike, str],
108    split: Literal["train", "test"],
109    grader: Literal["grader1", "grader2", "grader3"] = "grader1",
110    download: bool = False,
111) -> Tuple[List[str], List[str]]:
112    """Get paths to the CSC Fundus Segmentation data.
113
114    Args:
115        path: Filepath to a folder where the data is downloaded for further processing.
116        split: The choice of data split.
117        grader: The choice of annotator whose segmentations are used.
118        download: Whether to download the data if it is not present.
119
120    Returns:
121        List of filepaths for the image data.
122        List of filepaths for the label data.
123    """
124    assert split in ["train", "test"], f"'{split}' is not a valid split."
125
126    data_dir = get_csc_fundus_segmentation_data(path=path, grader=grader, download=download)
127
128    paired_paths = natsorted(glob(os.path.join(data_dir, split, "*.jpg")))
129    assert len(paired_paths) > 0, f"No images found for split '{split}' and grader '{grader}'."
130
131    image_dir = os.path.join(data_dir, "preprocessed", split, "images")
132    gt_dir = os.path.join(data_dir, "preprocessed", split, "masks")
133
134    image_paths, gt_paths = [], []
135    for paired_path in paired_paths:
136        image_path, gt_path = _split_paired_image(paired_path, image_dir, gt_dir)
137        image_paths.append(image_path)
138        gt_paths.append(gt_path)
139
140    return image_paths, gt_paths
141
142
143def get_csc_fundus_segmentation_dataset(
144    path: Union[os.PathLike, str],
145    patch_shape: Tuple[int, int],
146    split: Literal["train", "test"],
147    grader: Literal["grader1", "grader2", "grader3"] = "grader1",
148    resize_inputs: bool = False,
149    download: bool = False,
150    **kwargs
151) -> Dataset:
152    """Get the CSC Fundus Segmentation dataset for subretinal fluid segmentation in fundus photographs.
153
154    Args:
155        path: Filepath to a folder where the downloaded data will be saved.
156        patch_shape: The patch shape to use for training.
157        split: The choice of data split.
158        grader: The choice of annotator whose segmentations are used.
159        resize_inputs: Whether to resize the inputs to the expected patch shape.
160        download: Whether to download the data if it is not present.
161        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
162
163    Returns:
164        The segmentation dataset.
165    """
166    image_paths, gt_paths = get_csc_fundus_segmentation_paths(
167        path=path, split=split, grader=grader, download=download
168    )
169
170    if resize_inputs:
171        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
172        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
173            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs,
174        )
175
176    return torch_em.default_segmentation_dataset(
177        raw_paths=image_paths,
178        raw_key=None,
179        label_paths=gt_paths,
180        label_key=None,
181        patch_shape=patch_shape,
182        is_seg_dataset=False,
183        **kwargs
184    )
185
186
187def get_csc_fundus_segmentation_loader(
188    path: Union[os.PathLike, str],
189    batch_size: int,
190    patch_shape: Tuple[int, int],
191    split: Literal["train", "test"],
192    grader: Literal["grader1", "grader2", "grader3"] = "grader1",
193    resize_inputs: bool = False,
194    download: bool = False,
195    **kwargs
196) -> DataLoader:
197    """Get the CSC Fundus Segmentation dataloader for subretinal fluid segmentation in fundus photographs.
198
199    Args:
200        path: Filepath to a folder where the downloaded data will be saved.
201        batch_size: The batch size for training.
202        patch_shape: The patch shape to use for training.
203        split: The choice of data split.
204        grader: The choice of annotator whose segmentations are used.
205        resize_inputs: Whether to resize the inputs to the expected patch shape.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
208
209    Returns:
210        The DataLoader.
211    """
212    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
213    dataset = get_csc_fundus_segmentation_dataset(
214        path, patch_shape, split, grader, resize_inputs, download, **ds_kwargs
215    )
216    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)