torch_em.data.datasets.medical.etis_larib

  1"""The ETIS-Larib dataset contains annotations for polyp segmentation in colonoscopy images.
  2
  3The dataset consists of 196 still colonoscopy frames, each paired with a binary segmentation
  4mask of the polyp region.
  5
  6The dataset is located at https://www.kaggle.com/datasets/nguyenvoquocduong/etis-laribpolypdb. This
  7is a mirror of the original ETIS-Larib Polyp DB release, which is otherwise gated behind manual
  8registration on the same site family as CVC-ColonDB (https://polyp.grand-challenge.org).
  9
 10The dataset is from the publication https://doi.org/10.1007/s11548-013-0926-3.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 17from typing import Union, Tuple, List
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26KAGGLE_DATASET_NAME = "nguyenvoquocduong/etis-laribpolypdb"
 27
 28
 29def get_etis_larib_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the ETIS-Larib dataset.
 31
 32    Args:
 33        path: Filepath to a folder where the data is downloaded for further processing.
 34        download: Whether to download the data if it is not present.
 35
 36    Returns:
 37        Filepath where the data is downloaded.
 38    """
 39    data_dir = os.path.join(path, "images")
 40    if os.path.exists(data_dir):
 41        return path
 42
 43    os.makedirs(path, exist_ok=True)
 44
 45    util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 46
 47    zip_path = os.path.join(path, "etis-laribpolypdb.zip")
 48    util.unzip(zip_path=zip_path, dst=path)
 49
 50    if not os.path.exists(data_dir):
 51        raise RuntimeError(f"The dataset could not be found at '{path}' after extraction.")
 52
 53    return path
 54
 55
 56def get_etis_larib_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 57    """Get paths to the ETIS-Larib 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_etis_larib_data(path=path, download=download)
 68
 69    image_paths = natsorted(glob(os.path.join(data_dir, "images", "*.png")))
 70    gt_paths = natsorted(glob(os.path.join(data_dir, "masks", "*.png")))
 71
 72    if len(image_paths) == 0 or len(image_paths) != len(gt_paths):
 73        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 74
 75    return image_paths, gt_paths
 76
 77
 78def get_etis_larib_dataset(
 79    path: Union[os.PathLike, str],
 80    patch_shape: Tuple[int, int],
 81    resize_inputs: bool = False,
 82    download: bool = False,
 83    **kwargs
 84) -> Dataset:
 85    """Get the ETIS-Larib dataset for polyp segmentation in colonoscopy images.
 86
 87    Args:
 88        path: Filepath to a folder where the data is downloaded for further processing.
 89        patch_shape: The patch shape to use for training.
 90        resize_inputs: Whether to resize the inputs to the expected patch shape.
 91        download: Whether to download the data if it is not present.
 92        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 93
 94    Returns:
 95        The segmentation dataset.
 96    """
 97    image_paths, gt_paths = get_etis_larib_paths(path, download)
 98
 99    if resize_inputs:
100        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
101        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
102            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
103        )
104
105    return torch_em.default_segmentation_dataset(
106        raw_paths=image_paths,
107        raw_key=None,
108        label_paths=gt_paths,
109        label_key=None,
110        patch_shape=patch_shape,
111        is_seg_dataset=False,
112        **kwargs
113    )
114
115
116def get_etis_larib_loader(
117    path: Union[os.PathLike, str],
118    batch_size: int,
119    patch_shape: Tuple[int, int],
120    resize_inputs: bool = False,
121    download: bool = False,
122    **kwargs
123) -> DataLoader:
124    """Get the ETIS-Larib dataloader for polyp segmentation in colonoscopy images.
125
126    Args:
127        path: Filepath to a folder where the data is downloaded for further processing.
128        batch_size: The batch size for training.
129        patch_shape: The patch shape to use for training.
130        resize_inputs: Whether to resize the inputs to the expected patch shape.
131        download: Whether to download the data if it is not present.
132        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
133
134    Returns:
135        The DataLoader.
136    """
137    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
138    dataset = get_etis_larib_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
139    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)