torch_em.data.datasets.medical.ravir

  1"""The RAVIR dataset contains annotations for segmentation of retinal arteries and veins
  2in infrared reflectance imaging.
  3
  4The dataset is from the RAVIR challenge: https://ravir.grand-challenge.org/RAVIR/.
  5This dataset is from the publication https://doi.org/10.1109/JBHI.2022.3163352.
  6Please cite them if you use this dataset for your research.
  7"""
  8
  9import os
 10import shutil
 11from glob import glob
 12from typing import Union, Tuple, List
 13
 14import torch_em
 15
 16from .. import util
 17
 18
 19URL = "https://drive.google.com/uc?export=download&id=1ZlZoSStvE9VCRq3bJiGhQH931EF0h3hh"
 20CHECKSUM = "b9cc2e84660ab4ebeb583d510bd71057faf596a99ed6d1e27aee361e3a3f1381"
 21
 22
 23def get_ravir_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 24    """Download the RAVIR dataset.
 25
 26    Args:
 27        path: Filepath to a folder where the data is downloaded for further processing.
 28        download: Whether to download the data if it is not present.
 29
 30    Returns:
 31        Filepath where the data is downloaded.
 32    """
 33    data_dir = os.path.join(path, "RAVIR_Dataset")
 34    if os.path.exists(data_dir):
 35        return data_dir
 36
 37    os.makedirs(path, exist_ok=True)
 38
 39    zip_path = os.path.join(path, "ravir.zip")
 40    util.download_source_gdrive(
 41        path=zip_path, url=URL, download=download, checksum=CHECKSUM, download_type="zip"
 42    )
 43    util.unzip(zip_path=zip_path, dst=path)
 44
 45    # Updating the folder structure.
 46    tmp_dir = os.path.join(path, r"RAVIR Dataset")
 47    assert os.path.exists(tmp_dir), "Something went wrong with the data download"
 48    shutil.move(tmp_dir, data_dir)
 49
 50    return data_dir
 51
 52
 53def get_ravir_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[int], List[int]]:
 54    """Get paths to the RAVIR dataset.
 55
 56    Args:
 57        path: Filepath to a folder where the data is downloaded for further processing.
 58        download: Whether to download the data if it is not present.
 59
 60    Returns:
 61        List of filepaths for the image data.
 62        List of filepaths for the label data.
 63    """
 64    data_dir = get_ravir_data(path=path, download=download)
 65
 66    image_paths = sorted(glob(os.path.join(data_dir, "train", "training_images", "*")))
 67    gt_paths = sorted(glob(os.path.join(data_dir, "train", "training_masks", "*")))
 68
 69    return image_paths, gt_paths
 70
 71
 72def get_ravir_dataset(
 73    path: Union[os.PathLike, str],
 74    patch_shape: Tuple[int, int],
 75    resize_inputs: bool = False,
 76    download: bool = False,
 77    **kwargs
 78):
 79    """Get the RAVIR dataset for segmentation of retinal arteries and veins.
 80
 81    Args:
 82        path: Filepath to a folder where the data is downloaded for further processing.
 83        patch_shape: The patch shape to use for training.
 84        resize_inputs: Whether to resize the inputs to the patch shape.
 85        download: Whether to download the data if it is not present.
 86        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 87
 88    Returns:
 89        The segmentation dataset.
 90    """
 91    image_paths, gt_paths = get_ravir_paths(path, download)
 92
 93    if resize_inputs:
 94        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
 95        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
 96            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
 97        )
 98
 99    return torch_em.default_segmentation_dataset(
100        raw_paths=image_paths,
101        raw_key=None,
102        label_paths=gt_paths,
103        label_key=None,
104        patch_shape=patch_shape,
105        is_seg_dataset=False,
106        **kwargs
107    )
108
109
110def get_ravir_loader(
111    path: Union[os.PathLike, str],
112    batch_size: int,
113    patch_shape: Tuple[int, int],
114    resize_inputs: bool = False,
115    download: bool = False,
116    **kwargs
117):
118    """Get the RAVIR dataloader for segmentation of retinal arteries and veins.
119
120    Args:
121        path: Filepath to a folder where the data is downloaded for further processing.
122        batch_size: The batch size for training.
123        patch_shape: The patch shape to use for training.
124        download: Whether to download the data if it is not present.
125        resize_inputs: Whether to resize the inputs to the patch shape.
126        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
127
128    Returns:
129        The DataLoader.
130    """
131    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
132    dataset = get_ravir_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
133    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)