torch_em.data.datasets.medical.edd2020

  1"""The EDD2020 dataset contains annotations for multi-class disease segmentation in
  2gastrointestinal endoscopy images.
  3
  4The dataset consists of 386 endoscopy frames from 5 different institutions and multiple
  5GI organs (esophagus, stomach, colon), collected with white light, narrow-band imaging and
  6chromoendoscopy. Each image is annotated with pixel-level masks for up to 5 disease classes:
  7non-dysplastic Barrett's esophagus (BE), suspicious lesions, high-grade dysplasia (HGD),
  8cancer, and polyp. The official challenge data is gated behind manual registration on
  9https://edd2020.grand-challenge.org, so we instead use the openly mirrored copy at
 10https://www.kaggle.com/datasets/orvile/edd2020-endoscopy-detection-and-segmentation, which
 11matches the official release (386 images, same organizers and class structure). NOTE: the
 12original data release states the license as CC BY-NC-SA 4.0, while the Kaggle mirror lists it
 13as CC BY 4.0; please check the current license terms before using this data.
 14
 15This dataset is from the publication https://doi.org/10.1016/j.media.2021.102002.
 16Please cite it if you use this dataset for your research.
 17"""
 18
 19import os
 20from glob import glob
 21from tqdm import tqdm
 22from pathlib import Path
 23from natsort import natsorted
 24from typing import Union, Tuple, List
 25
 26import numpy as np
 27import imageio.v3 as imageio
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36KAGGLE_DATASET_NAME = "orvile/edd2020-endoscopy-detection-and-segmentation"
 37
 38CLASS_NAMES = ["BE", "suspicious", "HGD", "cancer", "polyp"]
 39
 40
 41def get_edd2020_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 42    """Download the EDD2020 dataset.
 43
 44    Args:
 45        path: Filepath to a folder where the data is downloaded for further processing.
 46        download: Whether to download the data if it is not present.
 47
 48    Returns:
 49        Filepath where the data is downloaded.
 50    """
 51    data_dir = os.path.join(path, "EDD2020")
 52    if os.path.exists(data_dir):
 53        return data_dir
 54
 55    os.makedirs(path, exist_ok=True)
 56
 57    util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 58
 59    zip_path = os.path.join(path, "edd2020-endoscopy-detection-and-segmentation.zip")
 60    util.unzip(zip_path=zip_path, dst=path)
 61
 62    if not os.path.exists(data_dir):
 63        raise RuntimeError(f"The dataset could not be found at '{path}' after extraction.")
 64
 65    return data_dir
 66
 67
 68def get_edd2020_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 69    """Get paths to the EDD2020 data.
 70
 71    Args:
 72        path: Filepath to a folder where the data is downloaded for further processing.
 73        download: Whether to download the data if it is not present.
 74
 75    Returns:
 76        List of filepaths for the image data.
 77        List of filepaths for the label data.
 78    """
 79    data_dir = get_edd2020_data(path=path, download=download)
 80
 81    image_paths = natsorted(glob(os.path.join(data_dir, "originalImages", "*.jpg")))
 82    if len(image_paths) == 0:
 83        raise RuntimeError("Something went wrong with fetching the image paths.")
 84
 85    neu_gt_dir = os.path.join(data_dir, "masks", "preprocessed")
 86    os.makedirs(neu_gt_dir, exist_ok=True)
 87
 88    gt_paths = []
 89    for image_path in tqdm(image_paths, desc="Preprocessing labels"):
 90        image_id = Path(image_path).stem
 91        gt_path = os.path.join(neu_gt_dir, f"{image_id}.tif")
 92        gt_paths.append(gt_path)
 93        if os.path.exists(gt_path):
 94            continue
 95
 96        image = imageio.imread(image_path)
 97        label = np.zeros(image.shape[:2], dtype="uint8")
 98        for class_id, class_name in enumerate(CLASS_NAMES, start=1):
 99            mask_path = os.path.join(data_dir, "masks", f"{image_id}_{class_name}.tif")
100            if not os.path.exists(mask_path):
101                continue
102            mask = imageio.imread(mask_path)
103            label[mask > 0] = class_id
104
105        imageio.imwrite(gt_path, label, compression="zlib")
106
107    return image_paths, gt_paths
108
109
110def get_edd2020_dataset(
111    path: Union[os.PathLike, str],
112    patch_shape: Tuple[int, int],
113    resize_inputs: bool = False,
114    download: bool = False,
115    **kwargs
116) -> Dataset:
117    """Get the EDD2020 dataset for multi-class disease segmentation.
118
119    Args:
120        path: Filepath to a folder where the data is downloaded for further processing.
121        patch_shape: The patch shape to use for training.
122        resize_inputs: Whether to resize the inputs to the patch shape.
123        download: Whether to download the data if it is not present.
124        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
125
126    Returns:
127        The segmentation dataset.
128    """
129    image_paths, gt_paths = get_edd2020_paths(path, download)
130
131    if resize_inputs:
132        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
133        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
134            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
135        )
136
137    return torch_em.default_segmentation_dataset(
138        raw_paths=image_paths,
139        raw_key=None,
140        label_paths=gt_paths,
141        label_key=None,
142        patch_shape=patch_shape,
143        is_seg_dataset=False,
144        **kwargs
145    )
146
147
148def get_edd2020_loader(
149    path: Union[os.PathLike, str],
150    batch_size: int,
151    patch_shape: Tuple[int, int],
152    resize_inputs: bool = False,
153    download: bool = False,
154    **kwargs
155) -> DataLoader:
156    """Get the EDD2020 dataloader for multi-class disease segmentation.
157
158    Args:
159        path: Filepath to a folder where the data is downloaded for further processing.
160        batch_size: The batch size for training.
161        patch_shape: The patch shape to use for training.
162        resize_inputs: Whether to resize the inputs to the patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
165
166    Returns:
167        The DataLoader.
168    """
169    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
170    dataset = get_edd2020_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
171    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)