torch_em.data.datasets.medical.rvo_me

  1"""RVO-ME (also released as 'RVO-Lesion') is a dataset for segmentation of macular lesions in
  2retinal vein occlusion (RVO) in optical coherence tomography (OCT) B-scans.
  3
  4The dataset consists of 3,012 OCT B-scan images from 130 patients (146 eyes). The pixel-level masks
  5label 4 classes: 0 = background, 1 = SRF (subretinal fluid), 2 = IRF (intraretinal fluid),
  63 = ELM (external limiting membrane), 4 = EZ (ellipsoid zone).
  7
  8This dataset is located at https://doi.org/10.6084/m9.figshare.29804435.v1 (figshare, CC BY 4.0).
  9The dataset is from the publication https://doi.org/10.1038/s41597-026-06695-5.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14from natsort import natsorted
 15from typing import Union, Tuple, Literal, List
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URL = "https://ndownloader.figshare.com/files/56848025"
 25CHECKSUM = "ffe522f1b09e1a4c0c8eae10f776b25e3d3e13a1da30755f96f8356c98af3776"
 26
 27
 28def get_rvo_me_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 29    """Download the RVO-ME data.
 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, "RVO-Lesion")
 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, "RVO-Lesion.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
 51def get_rvo_me_paths(
 52    path: Union[os.PathLike, str], split: Literal['train', 'test'] = "train", download: bool = False
 53) -> Tuple[List[str], List[str]]:
 54    """Get paths to the RVO-ME data.
 55
 56    Args:
 57        path: Filepath to a folder where the data is downloaded for further processing.
 58        split: The choice of data split. Either 'train' or 'test'.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        List of filepaths for the image data.
 63        List of filepaths for the label data.
 64    """
 65    data_dir = get_rvo_me_data(path, download)
 66
 67    if split not in ("train", "test"):
 68        raise ValueError(f"'{split}' is not a valid split. Choose either 'train' or 'test'.")
 69
 70    split_file = os.path.join(data_dir, "Image_Seg", f"{split}.txt")
 71    with open(split_file) as f:
 72        fnames = [line.strip() for line in f if line.strip()]
 73
 74    raw_paths = natsorted(os.path.join(data_dir, "Image_Seg", "images", fname) for fname in fnames)
 75    label_paths = natsorted(
 76        os.path.join(data_dir, "Image_Seg", "masks", os.path.splitext(fname)[0] + ".png") for fname in fnames
 77    )
 78
 79    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 80    assert all(os.path.exists(p) for p in raw_paths)
 81    assert all(os.path.exists(p) for p in label_paths)
 82
 83    return raw_paths, label_paths
 84
 85
 86def get_rvo_me_dataset(
 87    path: Union[os.PathLike, str],
 88    patch_shape: Tuple[int, int],
 89    split: Literal['train', 'test'] = "train",
 90    resize_inputs: bool = False,
 91    download: bool = False,
 92    **kwargs
 93) -> Dataset:
 94    """Get the RVO-ME dataset for segmentation of macular lesions in OCT B-scans.
 95
 96    Args:
 97        path: Filepath to a folder where the data is downloaded for further processing.
 98        patch_shape: The patch shape to use for training.
 99        split: The choice of data split. Either 'train' or 'test'.
100        resize_inputs: Whether to resize the inputs to the patch shape.
101        download: Whether to download the data if it is not present.
102        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
103
104    Returns:
105        The segmentation dataset.
106    """
107    raw_paths, label_paths = get_rvo_me_paths(path, split, download)
108
109    if resize_inputs:
110        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
111        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
112            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
113        )
114
115    return torch_em.default_segmentation_dataset(
116        raw_paths=raw_paths,
117        raw_key=None,
118        label_paths=label_paths,
119        label_key=None,
120        patch_shape=patch_shape,
121        is_seg_dataset=False,
122        **kwargs
123    )
124
125
126def get_rvo_me_loader(
127    path: Union[os.PathLike, str],
128    batch_size: int,
129    patch_shape: Tuple[int, int],
130    split: Literal['train', 'test'] = "train",
131    resize_inputs: bool = False,
132    download: bool = False,
133    **kwargs
134) -> DataLoader:
135    """Get the RVO-ME dataloader for segmentation of macular lesions in OCT B-scans.
136
137    Args:
138        path: Filepath to a folder where the data is downloaded for further processing.
139        batch_size: The batch size for training.
140        patch_shape: The patch shape to use for training.
141        split: The choice of data split. Either 'train' or 'test'.
142        resize_inputs: Whether to resize the inputs to the patch shape.
143        download: Whether to download the data if it is not present.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
145
146    Returns:
147        The DataLoader.
148    """
149    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
150    dataset = get_rvo_me_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
151    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)