torch_em.data.datasets.medical.polypgen

  1"""The PolypGen dataset contains annotations for polyp detection and segmentation in
  2colonoscopy images and video sequence frames, collected from six different medical
  3centres across Europe and Africa.
  4
  5NOTE: The dataset is hosted on Synapse. Downloading it requires the 'synapseclient'
  6python library and a Synapse account with an authentication token stored in the
  7'~/.synapseConfig' file. See 'get_polypgen_data' for details.
  8
  9The Synapse project 'syn26376615' has two versions of the multi-centre archive
 10('PolypGen2021_MultiCenterData_v2.zip' and '..._v3.zip'); this module downloads both
 11(via 'download_source_synapse') but only extracts and uses the newer v3 archive, which
 12unpacks into a 'data_C<i>/images_C<i>' and 'data_C<i>/masks_C<i>' folder per centre
 13(confirmed by inspecting the real archive contents, not assumed).
 14
 15The dataset is from the publication https://doi.org/10.1038/s41597-023-01981-y.
 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, Optional, Literal
 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
 36ENTITY = "syn26376615"
 37
 38CENTERS = ["C1", "C2", "C3", "C4", "C5", "C6"]
 39
 40
 41def get_polypgen_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 42    """Download the PolypGen dataset.
 43
 44    The dataset is located at https://www.synapse.org/Synapse:syn26376615.
 45
 46    Follow the instructions below to get access to the dataset.
 47    - Create a free account at https://www.synapse.org.
 48    - Generate a personal access token and store it in a '~/.synapseConfig' file, see
 49      https://python-docs.synapse.org/tutorials/authentication/ for details.
 50    - Install the 'synapseclient' python library.
 51
 52    Args:
 53        path: Filepath to a folder where the data is downloaded for further processing.
 54        download: Whether to download the data if it is not present.
 55
 56    Returns:
 57        Filepath where the data is downloaded.
 58    """
 59    extracted_dir = os.path.join(path, "PolypGen2021_MultiCenterData_v3")
 60    if os.path.exists(extracted_dir):
 61        return path
 62
 63    os.makedirs(path, exist_ok=True)
 64    util.download_source_synapse(path=path, entity=ENTITY, download=download)
 65
 66    zip_path = os.path.join(path, "PolypGen2021_MultiCenterData_v3.zip")
 67    util.unzip(zip_path=zip_path, dst=path, remove=False)
 68
 69    return path
 70
 71
 72def get_polypgen_paths(
 73    path: Union[os.PathLike, str],
 74    center: Optional[Literal["C1", "C2", "C3", "C4", "C5", "C6"]] = None,
 75    download: bool = False,
 76) -> Tuple[List[str], List[str]]:
 77    """Get paths to the PolypGen data.
 78
 79    Args:
 80        path: Filepath to a folder where the data is downloaded for further processing.
 81        center: The choice of a specific centre's data. By default, loads data from all centres.
 82        download: Whether to download the data if it is not present.
 83
 84    Returns:
 85        List of filepaths for the image data.
 86        List of filepaths for the label data.
 87    """
 88    data_dir = get_polypgen_data(path, download)
 89
 90    centers = CENTERS if center is None else [center]
 91
 92    image_paths, gt_paths = [], []
 93    for this_center in centers:
 94        this_image_paths = natsorted(
 95            glob(os.path.join(data_dir, "**", f"images_{this_center}", "*.jpg"), recursive=True)
 96        )
 97        for image_path in this_image_paths:
 98            # Masks live in a sibling 'masks_C<i>' folder and their filename has an extra '_mask' suffix,
 99            # e.g. 'images_C1/100H0050.jpg' pairs with 'masks_C1/100H0050_mask.jpg'.
100            mask_dir = os.path.dirname(image_path).replace(f"images_{this_center}", f"masks_{this_center}")
101            gt_path = os.path.join(mask_dir, f"{Path(image_path).stem}_mask.jpg")
102            if not os.path.exists(gt_path):
103                continue
104
105            image_paths.append(image_path)
106            gt_paths.append(gt_path)
107
108    # The masks are lossily JPEG-compressed grayscale images replicated across 3 channels (background
109    # near 0, foreground near 255), not the single-channel binary masks 'ImageCollectionDataset' expects.
110    # They are binarized once here and cached as '.tif' files next to the original masks.
111    neu_gt_paths = []
112    for gt_path in tqdm(gt_paths, desc="Preprocessing PolypGen masks"):
113        neu_gt_path = os.path.join(os.path.dirname(gt_path), f"{Path(gt_path).stem}.tif")
114        neu_gt_paths.append(neu_gt_path)
115        if os.path.exists(neu_gt_path):
116            continue
117
118        gt = imageio.imread(gt_path)
119        if gt.ndim == 3:  # Some centres' masks are single-channel grayscale rather than RGB.
120            gt = np.mean(gt, axis=-1)
121        gt = (gt > 128).astype("uint8")
122        imageio.imwrite(neu_gt_path, gt, compression="zlib")
123    gt_paths = neu_gt_paths
124
125    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0, (
126        "No image-mask pairs were found. The expected per-centre 'images_C<i>' / 'masks_C<i>' folder layout "
127        "may not match the actual structure of the downloaded PolypGen data. Please inspect the data at "
128        f"'{data_dir}' and update the search pattern in 'get_polypgen_paths' accordingly."
129    )
130
131    return image_paths, gt_paths
132
133
134def get_polypgen_dataset(
135    path: Union[os.PathLike, str],
136    patch_shape: Tuple[int, int],
137    center: Optional[Literal["C1", "C2", "C3", "C4", "C5", "C6"]] = None,
138    resize_inputs: bool = False,
139    download: bool = False,
140    **kwargs
141) -> Dataset:
142    """Get the PolypGen dataset for polyp segmentation in colonoscopy images.
143
144    Args:
145        path: Filepath to a folder where the data is downloaded for further processing.
146        patch_shape: The patch shape to use for training.
147        center: The choice of a specific centre's data. By default, loads data from all centres.
148        resize_inputs: Whether to resize inputs to the desired patch shape.
149        download: Whether to download the data if it is not present.
150        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
151
152    Returns:
153        The segmentation dataset.
154    """
155    image_paths, gt_paths = get_polypgen_paths(path, center, download)
156
157    if resize_inputs:
158        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
159        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
160            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
161        )
162
163    return torch_em.default_segmentation_dataset(
164        raw_paths=image_paths,
165        raw_key=None,
166        label_paths=gt_paths,
167        label_key=None,
168        patch_shape=patch_shape,
169        is_seg_dataset=False,
170        **kwargs
171    )
172
173
174def get_polypgen_loader(
175    path: Union[os.PathLike, str],
176    batch_size: int,
177    patch_shape: Tuple[int, int],
178    center: Optional[Literal["C1", "C2", "C3", "C4", "C5", "C6"]] = None,
179    resize_inputs: bool = False,
180    download: bool = False,
181    **kwargs
182) -> DataLoader:
183    """Get the PolypGen dataloader for polyp segmentation in colonoscopy images.
184
185    Args:
186        path: Filepath to a folder where the data is downloaded for further processing.
187        batch_size: The batch size for training.
188        patch_shape: The patch shape to use for training.
189        center: The choice of a specific centre's data. By default, loads data from all centres.
190        resize_inputs: Whether to resize inputs to the desired patch shape.
191        download: Whether to download the data if it is not present.
192        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
193
194    Returns:
195        The DataLoader.
196    """
197    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
198    dataset = get_polypgen_dataset(path, patch_shape, center, resize_inputs, download, **ds_kwargs)
199    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)