torch_em.data.datasets.light_microscopy.apacs23

  1"""The APACS23 dataset contains annotations for cell segmentation in
  2digitized Pap smear images of the cervix.
  3
  4APACS23 stands for Annotated PAp smear images for Cell Segmentation 2023. The dataset holds 3535
  5image and mask pairs of 2000 x 2000 pixels, which the authors split into a training and a test part.
  6
  7NOTE: The masks are binary. They separate the cells from the background, and they hold no instance
  8id and no class. The publication reports about 37000 segmented cells, because the annotators
  9outlined every cell by hand, but the released masks merge them into one foreground.
 10
 11NOTE: The repository stores every file on its own, and it offers no archive. The loader therefore
 12downloads about 7000 files, which takes a while on the first call. It stores them under `path`, so
 13a later call reads them from disk.
 14
 15NOTE: A few files have no partner. The training part holds 10 images without a mask and 10 masks
 16without an image, and the test part holds 5 of each. The loader skips them, so it yields 2207
 17training pairs and 1328 test pairs.
 18
 19The dataset is located at https://doi.org/10.17605/OSF.IO/CKA2F under the CC BY 4.0 license.
 20This dataset is from the publication https://doi.org/10.1038/s41597-024-03566-9.
 21Please cite it if you use this dataset in your research.
 22"""
 23
 24import os
 25import time
 26import json
 27import urllib.error
 28import urllib.request
 29from glob import glob
 30from natsort import natsorted
 31from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
 32from typing import List, Literal, Tuple, Union
 33
 34from torch.utils.data import DataLoader, Dataset
 35
 36import torch_em
 37
 38from .. import util
 39
 40
 41OSF_NODE = "cka2f"
 42OSF_API = f"https://api.osf.io/v2/nodes/{OSF_NODE}/files/osfstorage/"
 43
 44# The repository limits how many requests it serves per minute, so the download waits between files.
 45REQUEST_DELAY = 0.3
 46RETRY_DELAY = 5.0
 47MAX_RETRIES = 6
 48
 49# The folder of a split, and the folders that hold its images and its masks.
 50SPLITS = {
 51    "train": ("training", "APACS23_Training_Input", "APACS23_Training_GroundTruth"),
 52    "test": ("test", "APACS23_Test_Input", "APACS23_Test_GroundTruth"),
 53}
 54
 55
 56def _list_folder(url: str) -> List[dict]:
 57    """List a folder of the repository, and follow its pages."""
 58    parts = urlparse(url)
 59    query = parse_qs(parts.query)
 60    query["page[size]"] = ["100"]
 61    url = urlunparse(parts._replace(query=urlencode(query, doseq=True)))
 62
 63    entries, seen = [], set()
 64    while url:
 65        with urllib.request.urlopen(url, timeout=120) as response:
 66            page = json.load(response)
 67        for entry in page["data"]:
 68            # A page can repeat an entry, so the id decides whether it is new.
 69            if entry["id"] not in seen:
 70                seen.add(entry["id"])
 71                entries.append(entry)
 72        url = page["links"].get("next")
 73    return entries
 74
 75
 76def _download_file(url: str, output_path: str) -> None:
 77    """Download one file, and wait when the repository refuses the request.
 78
 79    The repository answers with the status 403 once too many requests arrive in a short time, so
 80    every failed try waits longer than the one before it.
 81    """
 82    request = urllib.request.Request(url, headers={"User-Agent": "torch-em"})
 83    for attempt in range(MAX_RETRIES):
 84        try:
 85            with urllib.request.urlopen(request, timeout=120) as response:
 86                content = response.read()
 87            with open(output_path, "wb") as f:
 88                f.write(content)
 89            return
 90        except urllib.error.HTTPError as error:
 91            if error.code != 403 or attempt == MAX_RETRIES - 1:
 92                raise
 93            time.sleep(RETRY_DELAY * 2 ** attempt)
 94
 95    raise RuntimeError(f"Could not download {url}.")
 96
 97
 98def _download_folder(entry: dict, output_dir: str) -> None:
 99    """Download every file of one repository folder."""
100    from tqdm import tqdm
101
102    os.makedirs(output_dir, exist_ok=True)
103    files = _list_folder(entry["relationships"]["files"]["links"]["related"]["href"])
104
105    for item in tqdm(files, desc=f"Download '{os.path.basename(output_dir)}'"):
106        name = item["attributes"]["name"]
107        output_path = os.path.join(output_dir, name)
108        if os.path.exists(output_path):
109            continue
110        _download_file(item["links"]["download"], output_path)
111        time.sleep(REQUEST_DELAY)
112
113
114def get_apacs23_data(
115    path: Union[os.PathLike, str],
116    split: Literal["train", "test"] = "train",
117    download: bool = False,
118) -> Tuple[str, str]:
119    """Download the APACS23 dataset.
120
121    Args:
122        path: Filepath to a folder where the downloaded data will be saved.
123        split: The data split. Either 'train' or 'test'.
124        download: Whether to download the data if it is not present.
125
126    Returns:
127        The filepath to the folder with the images.
128        The filepath to the folder with the masks.
129    """
130    if split not in SPLITS:
131        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
132
133    folder_name, image_folder, label_folder = SPLITS[split]
134    image_dir = os.path.join(path, split, "images")
135    label_dir = os.path.join(path, split, "masks")
136    if os.path.exists(image_dir) and os.path.exists(label_dir):
137        return image_dir, label_dir
138
139    if not download:
140        raise RuntimeError(f"Cannot find the data at {os.path.join(path, split)}, but download was set to False.")
141
142    os.makedirs(path, exist_ok=True)
143    top = _list_folder(OSF_API)
144    split_entry = next((e for e in top if e["attributes"]["name"] == folder_name), None)
145    if split_entry is None:
146        raise RuntimeError(f"Could not find the folder '{folder_name}' in the APACS23 repository.")
147
148    inner = _list_folder(split_entry["relationships"]["files"]["links"]["related"]["href"])
149    for entry in inner:
150        name = entry["attributes"]["name"]
151        if name == image_folder:
152            _download_folder(entry, image_dir)
153        elif name == label_folder:
154            _download_folder(entry, label_dir)
155
156    return image_dir, label_dir
157
158
159def get_apacs23_paths(
160    path: Union[os.PathLike, str],
161    split: Literal["train", "test"] = "train",
162    download: bool = False,
163) -> Tuple[List[str], List[str]]:
164    """Get paths to the APACS23 data.
165
166    Args:
167        path: Filepath to a folder where the downloaded data will be saved.
168        split: The data split. Either 'train' or 'test'.
169        download: Whether to download the data if it is not present.
170
171    Returns:
172        List of filepaths for the image data.
173        List of filepaths for the label data.
174    """
175    image_dir, label_dir = get_apacs23_data(path, split, download)
176
177    image_paths, label_paths = [], []
178    for image_path in natsorted(glob(os.path.join(image_dir, "*.jpg"))):
179        stem = os.path.splitext(os.path.basename(image_path))[0]
180        label_path = os.path.join(label_dir, f"{stem}.png")
181        # A few images have no mask, and a few masks have no image.
182        if not os.path.exists(label_path):
183            continue
184        image_paths.append(image_path)
185        label_paths.append(label_path)
186
187    if not image_paths:
188        raise RuntimeError(f"Could not find any APACS23 data in {image_dir}.")
189
190    return image_paths, label_paths
191
192
193def get_apacs23_dataset(
194    path: Union[os.PathLike, str],
195    patch_shape: Tuple[int, int],
196    split: Literal["train", "test"] = "train",
197    download: bool = False,
198    **kwargs,
199) -> Dataset:
200    """Get the APACS23 dataset for cell segmentation.
201
202    Args:
203        path: Filepath to a folder where the downloaded data will be saved.
204        patch_shape: The 2D patch shape to use for training.
205        split: The data split. Either 'train' or 'test'.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
208
209    Returns:
210        The segmentation dataset.
211    """
212    if len(patch_shape) != 2:
213        raise ValueError(f"The APACS23 patch shape must be two-dimensional, got {patch_shape}.")
214
215    image_paths, label_paths = get_apacs23_paths(path, split, download)
216    kwargs = util.ensure_transforms(ndim=2, **kwargs)
217
218    return torch_em.default_segmentation_dataset(
219        raw_paths=image_paths,
220        raw_key=None,
221        label_paths=label_paths,
222        label_key=None,
223        patch_shape=patch_shape,
224        is_seg_dataset=False,
225        ndim=2,
226        **kwargs,
227    )
228
229
230def get_apacs23_loader(
231    path: Union[os.PathLike, str],
232    batch_size: int,
233    patch_shape: Tuple[int, int],
234    split: Literal["train", "test"] = "train",
235    download: bool = False,
236    **kwargs,
237) -> DataLoader:
238    """Get the APACS23 dataloader for cell segmentation.
239
240    Args:
241        path: Filepath to a folder where the downloaded data will be saved.
242        batch_size: The batch size for training.
243        patch_shape: The 2D patch shape to use for training.
244        split: The data split. Either 'train' or 'test'.
245        download: Whether to download the data if it is not present.
246        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
247
248    Returns:
249        The DataLoader.
250    """
251    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
252    dataset = get_apacs23_dataset(
253        path=path, patch_shape=patch_shape, split=split, download=download, **ds_kwargs,
254    )
255    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)