torch_em.data.datasets.medical.corneal_confocal_nerve

  1"""The Corneal Confocal Nerve dataset contains annotations for pixel-level corneal nerve segmentation
  2in in-vivo corneal confocal microscopy (CCM) images.
  3
  4The dataset comprises 410 images from 88 participants, collected across two independently acquired
  5subsets with distinct acquisition conditions and participant groups. Each image is paired with an
  6expert-reviewed pixel-level nerve mask, in contrast to most other publicly available CCM datasets, which
  7only provide centreline annotations.
  8
  9The dataset is located at https://doi.org/10.5281/zenodo.18779434 (CC BY 4.0).
 10
 11This dataset is from the publication https://doi.org/10.1038/s41597-026-07418-6.
 12Please cite it if you use this dataset in your research.
 13"""
 14
 15import os
 16from glob import glob
 17from natsort import natsorted
 18from typing import Union, Tuple, List
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://zenodo.org/records/18779434/files/Dataset.zip"
 28CHECKSUM = "34c657a250487db58c17b3de0407f1a38fc85923f6886f1a52c193c04c995595"
 29
 30
 31def get_corneal_confocal_nerve_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the Corneal Confocal Nerve dataset.
 33
 34    Args:
 35        path: Filepath to a folder where the data is downloaded for further processing.
 36        download: Whether to download the data if it is not present.
 37
 38    Returns:
 39        Filepath where the data is downloaded.
 40    """
 41    data_dir = os.path.join(path, "Dataset")
 42    if os.path.exists(data_dir):
 43        return data_dir
 44
 45    os.makedirs(path, exist_ok=True)
 46
 47    zip_path = os.path.join(path, "Dataset.zip")
 48    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 49    util.unzip(zip_path=zip_path, dst=path)
 50
 51    return data_dir
 52
 53
 54def get_corneal_confocal_nerve_paths(
 55    path: Union[os.PathLike, str], download: bool = False
 56) -> Tuple[List[str], List[str]]:
 57    """Get paths to the Corneal Confocal Nerve data.
 58
 59    Args:
 60        path: Filepath to a folder where the data is downloaded for further processing.
 61        download: Whether to download the data if it is not present.
 62
 63    Returns:
 64        List of filepaths for the image data.
 65        List of filepaths for the label data.
 66    """
 67    data_dir = get_corneal_confocal_nerve_data(path, download)
 68
 69    image_paths = natsorted(glob(os.path.join(data_dir, "images", "*.png")))
 70    label_paths = [p.replace(os.sep + "images" + os.sep, os.sep + "annotations" + os.sep) for p in image_paths]
 71
 72    assert len(image_paths) > 0, f"Could not find any images in '{data_dir}'."
 73    for label_path in label_paths:
 74        assert os.path.exists(label_path), label_path
 75
 76    return image_paths, label_paths
 77
 78
 79def get_corneal_confocal_nerve_dataset(
 80    path: Union[os.PathLike, str],
 81    patch_shape: Tuple[int, int],
 82    resize_inputs: bool = False,
 83    download: bool = False,
 84    **kwargs
 85) -> Dataset:
 86    """Get the Corneal Confocal Nerve dataset for nerve segmentation in CCM images.
 87
 88    Args:
 89        path: Filepath to a folder where the data is downloaded for further processing.
 90        patch_shape: The patch shape to use for training.
 91        resize_inputs: Whether to resize the inputs to the expected patch shape.
 92        download: Whether to download the data if it is not present.
 93        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 94
 95    Returns:
 96        The segmentation dataset.
 97    """
 98    image_paths, label_paths = get_corneal_confocal_nerve_paths(path, download)
 99
100    if resize_inputs:
101        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
102        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
103            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
104        )
105
106    return torch_em.default_segmentation_dataset(
107        raw_paths=image_paths,
108        raw_key=None,
109        label_paths=label_paths,
110        label_key=None,
111        patch_shape=patch_shape,
112        is_seg_dataset=False,
113        **kwargs
114    )
115
116
117def get_corneal_confocal_nerve_loader(
118    path: Union[os.PathLike, str],
119    batch_size: int,
120    patch_shape: Tuple[int, int],
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> DataLoader:
125    """Get the Corneal Confocal Nerve dataloader for nerve segmentation in CCM images.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        batch_size: The batch size for training.
130        patch_shape: The patch shape to use for training.
131        resize_inputs: Whether to resize the inputs to the expected patch shape.
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
134
135    Returns:
136        The DataLoader.
137    """
138    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
139    dataset = get_corneal_confocal_nerve_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
140    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)