torch_em.data.datasets.medical.sustech_sysu

  1"""The SUSTech-SYSU dataset contains 1219 color fundus images from diabetic retinopathy (DR)
  2patients and healthy controls. 564 of these images have exudate annotations, provided as
  3Pascal VOC style bounding boxes around hard ('ex') and soft ('se') exudate lesions. These
  4bounding boxes are rasterized into binary exudate masks here. The dataset also ships DR grades,
  5optic disc bounding boxes and fovea locations, which are not exposed by this loader.
  6
  7The dataset is located at https://doi.org/10.6084/m9.figshare.12570770.
  8This dataset is from the publication https://doi.org/10.1038/s41597-020-00755-0.
  9Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13import xml.etree.ElementTree as ET
 14from glob import glob
 15from natsort import natsorted
 16from typing import Union, Tuple, List
 17
 18import numpy as np
 19import imageio.v3 as imageio
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://ndownloader.figshare.com/files/25320596"
 29CHECKSUM = "b5e3f31f7fc26f612f5fc04fbc8137a023d4812306ce8d2f92b6011dacd52735"
 30
 31
 32def get_sustech_sysu_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the SUSTech-SYSU dataset.
 34
 35    Args:
 36        path: Filepath to a folder where the data is downloaded for further processing.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath where the data is downloaded.
 41    """
 42    data_dir = os.path.join(path, "originalImages")
 43    if os.path.exists(data_dir):
 44        return path
 45
 46    os.makedirs(path, exist_ok=True)
 47
 48    zip_path = os.path.join(path, "sustech_sysu.zip")
 49    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 50    util.unzip(zip_path=zip_path, dst=path)
 51
 52    return path
 53
 54
 55def _bndbox_to_mask(xml_path, mask_path):
 56    root = ET.parse(xml_path).getroot()
 57    size = root.find("size")
 58    height, width = int(size.find("height").text), int(size.find("width").text)
 59
 60    mask = np.zeros((height, width), dtype="uint8")
 61    for obj in root.findall("object"):
 62        bndbox = obj.find("bndbox")
 63        xmin, ymin = int(bndbox.find("xmin").text), int(bndbox.find("ymin").text)
 64        xmax, ymax = int(bndbox.find("xmax").text), int(bndbox.find("ymax").text)
 65        mask[ymin:ymax, xmin:xmax] = 1
 66
 67    imageio.imwrite(mask_path, mask)
 68
 69
 70def get_sustech_sysu_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 71    """Get paths to the SUSTech-SYSU data.
 72
 73    Args:
 74        path: Filepath to a folder where the data is downloaded for further processing.
 75        download: Whether to download the data if it is not present.
 76
 77    Returns:
 78        List of filepaths for the image data.
 79        List of filepaths for the label data.
 80    """
 81    data_dir = get_sustech_sysu_data(path=path, download=download)
 82
 83    mask_dir = os.path.join(data_dir, "exudatesMasks")
 84    os.makedirs(mask_dir, exist_ok=True)
 85
 86    xml_paths = natsorted(glob(os.path.join(data_dir, "exudatesLabels", "*.xml")))
 87
 88    image_paths, label_paths = [], []
 89    for xml_path in xml_paths:
 90        fname = os.path.splitext(os.path.basename(xml_path))[0]
 91        image_path = os.path.join(data_dir, "originalImages", f"{fname}.jpg")
 92        assert os.path.exists(image_path), f"The image at '{image_path}' does not exist."
 93
 94        mask_path = os.path.join(mask_dir, f"{fname}.tif")
 95        if not os.path.exists(mask_path):
 96            _bndbox_to_mask(xml_path, mask_path)
 97
 98        image_paths.append(image_path)
 99        label_paths.append(mask_path)
100
101    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
102
103    return image_paths, label_paths
104
105
106def get_sustech_sysu_dataset(
107    path: Union[os.PathLike, str],
108    patch_shape: Tuple[int, int],
109    resize_inputs: bool = False,
110    download: bool = False,
111    **kwargs
112) -> Dataset:
113    """Get the SUSTech-SYSU dataset for exudate segmentation in fundus images.
114
115    Args:
116        path: Filepath to a folder where the data is downloaded for further processing.
117        patch_shape: The patch shape to use for training.
118        resize_inputs: Whether to resize the inputs to the expected patch shape.
119        download: Whether to download the data if it is not present.
120        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
121
122    Returns:
123        The segmentation dataset.
124    """
125    image_paths, label_paths = get_sustech_sysu_paths(path=path, download=download)
126
127    if resize_inputs:
128        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
129        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
130            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
131        )
132
133    return torch_em.default_segmentation_dataset(
134        raw_paths=image_paths,
135        raw_key=None,
136        label_paths=label_paths,
137        label_key=None,
138        patch_shape=patch_shape,
139        is_seg_dataset=False,
140        **kwargs
141    )
142
143
144def get_sustech_sysu_loader(
145    path: Union[os.PathLike, str],
146    batch_size: int,
147    patch_shape: Tuple[int, int],
148    resize_inputs: bool = False,
149    download: bool = False,
150    **kwargs
151) -> DataLoader:
152    """Get the SUSTech-SYSU dataloader for exudate segmentation in fundus images.
153
154    Args:
155        path: Filepath to a folder where the data is downloaded for further processing.
156        batch_size: The batch size for training.
157        patch_shape: The patch shape to use for training.
158        resize_inputs: Whether to resize the inputs to the expected patch shape.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_sustech_sysu_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)