torch_em.data.datasets.medical.far_polyp_seg

  1"""FAR-POLYP-SEG is a prospectively collected colonoscopy dataset for colorectal polyp
  2segmentation, acquired at Farhikhtegan Hospital, Islamic Azad University, Tehran, Iran
  3(February-December 2025), with clinician-created and gastroenterologist-reviewed
  4pixel-level masks.
  5
  6The dataset contains 8,181 RGB frames from 455 patients: 432 polyp-positive frames with
  7expert-verified binary masks and 7,749 normal colonic mucosa frames (without polyps or
  8masks). This module exposes only the 432 polyp-positive image-mask pairs, as the
  9normal-mucosa frames have no corresponding annotations to train a segmentation model on.
 10
 11The dataset is located at https://doi.org/10.5281/zenodo.20284781 and is licensed under
 12CC-BY-4.0.
 13
 14This dataset is from the publication https://doi.org/10.1007/s10278-026-02268-5.
 15Please cite it if you use this dataset for your research.
 16"""
 17
 18import os
 19from glob import glob
 20from tqdm import tqdm
 21from pathlib import Path
 22from natsort import natsorted
 23from typing import Union, Tuple, List
 24
 25import numpy as np
 26import imageio.v3 as imageio
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://zenodo.org/records/20284781/files/Dataset.zip"
 36CHECKSUM = "026a7b1e2ce9407e77b73373e999afcdbcbac776c5ac84c7944acf6285ed994d"
 37
 38
 39def get_far_polyp_seg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 40    """Download the FAR-POLYP-SEG data.
 41
 42    Args:
 43        path: Filepath to a folder where the data is downloaded for further processing.
 44        download: Whether to download the data if it is not present.
 45
 46    Returns:
 47        Filepath where the data is downloaded.
 48    """
 49    data_dir = os.path.join(path, "Dataset")
 50    if os.path.exists(data_dir):
 51        return data_dir
 52
 53    os.makedirs(path, exist_ok=True)
 54
 55    zip_path = os.path.join(path, "Dataset.zip")
 56    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 57    util.unzip(zip_path=zip_path, dst=path)
 58
 59    return data_dir
 60
 61
 62def get_far_polyp_seg_paths(
 63    path: Union[os.PathLike, str], download: bool = False
 64) -> Tuple[List[str], List[str]]:
 65    """Get paths to the FAR-POLYP-SEG data.
 66
 67    Args:
 68        path: Filepath to a folder where the data is downloaded for further processing.
 69        download: Whether to download the data if it is not present.
 70
 71    Returns:
 72        List of filepaths for the image data.
 73        List of filepaths for the label data.
 74    """
 75    data_dir = get_far_polyp_seg_data(path, download)
 76
 77    image_paths = natsorted(glob(os.path.join(data_dir, "Patient*", "Polyp", "*.jpg")))
 78
 79    # The masks are lossily JPEG-compressed grayscale images (background near 0, foreground
 80    # near 255), not the clean binary masks 'ImageCollectionDataset' expects. They are
 81    # binarized once here and cached as '.tif' files next to the original masks.
 82    gt_paths = []
 83    for image_path in tqdm(image_paths, desc="Preprocessing FAR-POLYP-SEG masks"):
 84        patient_dir = os.path.dirname(os.path.dirname(image_path))
 85        mask_path = os.path.join(patient_dir, "BinaryMask", os.path.basename(image_path))
 86        neu_gt_path = os.path.join(patient_dir, "BinaryMask", f"{Path(mask_path).stem}.tif")
 87        gt_paths.append(neu_gt_path)
 88        if os.path.exists(neu_gt_path):
 89            continue
 90
 91        mask = imageio.imread(mask_path)
 92        if mask.ndim == 3:
 93            mask = np.mean(mask, axis=-1)
 94        mask = (mask > 128).astype("uint8")
 95        imageio.imwrite(neu_gt_path, mask, compression="zlib")
 96
 97    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0, (
 98        "No image-mask pairs were found. The expected per-patient 'Polyp' / 'BinaryMask' folder layout "
 99        "may not match the actual structure of the downloaded data. Please inspect the data at "
100        f"'{data_dir}'."
101    )
102
103    return image_paths, gt_paths
104
105
106def get_far_polyp_seg_dataset(
107    path: Union[os.PathLike, str],
108    patch_shape: Tuple[int, int],
109    resize_inputs: bool = False,
110    download: bool = False,
111    **kwargs
112) -> Dataset:
113    """Get the FAR-POLYP-SEG dataset for polyp segmentation in colonoscopy images.
114
115    Args:
116        path: Filepath to a folder where the data is downloaded for further processing.
117        patch_shape: The patch shape to use for training.
118        resize_inputs: Whether to resize inputs to the desired patch shape.
119        download: Whether to download the data if it is not present.
120        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
121
122    Returns:
123        The segmentation dataset.
124    """
125    image_paths, gt_paths = get_far_polyp_seg_paths(path, download)
126
127    if resize_inputs:
128        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
129        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
130            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
131        )
132
133    return torch_em.default_segmentation_dataset(
134        raw_paths=image_paths,
135        raw_key=None,
136        label_paths=gt_paths,
137        label_key=None,
138        patch_shape=patch_shape,
139        is_seg_dataset=False,
140        **kwargs
141    )
142
143
144def get_far_polyp_seg_loader(
145    path: Union[os.PathLike, str],
146    batch_size: int,
147    patch_shape: Tuple[int, int],
148    resize_inputs: bool = False,
149    download: bool = False,
150    **kwargs
151) -> DataLoader:
152    """Get the FAR-POLYP-SEG dataloader for polyp segmentation in colonoscopy images.
153
154    Args:
155        path: Filepath to a folder where the data is downloaded for further processing.
156        batch_size: The batch size for training.
157        patch_shape: The patch shape to use for training.
158        resize_inputs: Whether to resize inputs to the desired patch shape.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_far_polyp_seg_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)