torch_em.data.datasets.medical.cystods

  1"""The CystoDS dataset contains annotations for tumor, ureteral orifice, resection scar and air bubble
  2segmentation in white light cystoscopy images.
  3
  4The full dataset consists of 8,067 images from 160 patients, labelled with five classes and
  522 subclasses for bladder cancer detection. Of these, 768 images have additional pixel-level
  6segmentation annotations (stored as polygon shapes). This module exposes the 763 of these 768
  7images that are actually present in the 'images' folder of the OSF release (5 filenames are
  8listed with a segmentation in the metadata but are missing from the images folder, see
  9`MISSING_IMAGES`). The remaining, unsegmented images only have classification labels and are
 10out of scope for segmentation.
 11
 12The dataset is located at https://osf.io/xvdhy (CC BY 4.0).
 13This dataset is from the publication https://doi.org/10.1038/s41597-026-06887-z.
 14Please cite it if you use this dataset in your research.
 15"""
 16
 17import os
 18import csv
 19import json
 20from glob import glob
 21from warnings import warn
 22from tqdm import tqdm
 23from natsort import natsorted
 24from typing import List, Tuple, Union
 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
 36OSF_NODE_ID = "xvdhy"
 37CSV_URL = "https://osf.io/download/n8cxz/"
 38CSV_CHECKSUM = "7988766bf2d17e7607ab2fcd9ad58d20a5fa4d7e107838a3b33b779e4e423f13"
 39
 40IMAGES_FOLDER_ID = "6883c4851f60df395a84faee"
 41SEGMENTATIONS_FOLDER_ID = "688413af2a29f9bd26e7cc33"
 42
 43LABEL_MAP = {
 44    "Tumor": 1,
 45    "Flat Tumor": 2,
 46    "Left ureteral orifice": 3,
 47    "Right ureteral orifice": 4,
 48    "Resection scar": 5,
 49    "Air bubble": 6,
 50}
 51
 52# 5 of the 768 filenames listed with 'json' == 1 in 'cystods.csv' are not actually present in the
 53# 'images' folder on OSF. This is a minor inconsistency in the released dataset, so they are excluded here.
 54MISSING_IMAGES = {"0ba96ad9.png", "23a59431.png", "4336849b.png", "9699deab.png", "d412d37a.png"}
 55
 56
 57def _osf_list_all(url):
 58    import time
 59    import requests
 60
 61    items = []
 62    expected_total = None
 63    while url:
 64        for attempt in range(5):
 65            r = requests.get(url)
 66            if r.status_code < 500:
 67                break
 68            time.sleep(2 ** attempt)
 69        r.raise_for_status()
 70        payload = r.json()
 71        items.extend(payload["data"])
 72        if expected_total is None:
 73            expected_total = payload["links"]["meta"]["total"]
 74        url = payload["links"].get("next")
 75    return items, expected_total
 76
 77
 78def _get_osf_folder_download_urls(folder_id, attempts=3):
 79    url = f"https://api.osf.io/v2/nodes/{OSF_NODE_ID}/files/osfstorage/{folder_id}/?page[size]=100"
 80    for attempt in range(attempts):
 81        items, expected_total = _osf_list_all(url)
 82        if len(items) == expected_total:
 83            return {item["attributes"]["name"]: item["links"]["download"] for item in items}
 84        warn(
 85            f"Listing the OSF folder {folder_id} returned {len(items)} items, expected {expected_total}. "
 86            "Retrying." if attempt < attempts - 1 else "Giving up."
 87        )
 88    return {item["attributes"]["name"]: item["links"]["download"] for item in items}
 89
 90
 91def _download_with_retries(path, url, download, attempts=5):
 92    import time
 93
 94    for attempt in range(attempts):
 95        try:
 96            util.download_source(path=path, url=url, download=download)
 97            return
 98        except Exception as e:
 99            if attempt == attempts - 1:
100                raise
101            warn(f"Download of {url} failed ({e}), retrying.")
102            time.sleep(2 ** attempt)
103
104
105def _get_segmented_filenames(csv_path):
106    with open(csv_path, encoding="utf-8-sig") as f:
107        rows = list(csv.DictReader(f))
108    filenames = [row["filename"] for row in rows if row.get("json") == "1"]
109    return natsorted(fname for fname in filenames if fname not in MISSING_IMAGES)
110
111
112def _rasterize_labels(segmentation_dir, filenames, label_dir):
113    if os.path.exists(label_dir) and len(glob(os.path.join(label_dir, "*.tif"))) == len(filenames):
114        return
115
116    os.makedirs(label_dir, exist_ok=True)
117
118    from skimage.draw import polygon as draw_polygon
119
120    for fname in tqdm(filenames, desc="Rasterize the CystoDS annotations"):
121        name = os.path.splitext(fname)[0]
122        label_path = os.path.join(label_dir, f"{name}.tif")
123        if os.path.exists(label_path):
124            continue
125
126        with open(os.path.join(segmentation_dir, f"{name}.json")) as f:
127            annotation = json.load(f)
128
129        image = imageio.imread(os.path.join(os.path.dirname(label_dir), "images", fname))
130        shape = image.shape[:2]
131
132        labels = np.zeros(shape, dtype="uint8")
133        for shape_annotation in annotation["shapes"]:
134            label_id = LABEL_MAP[shape_annotation["label"]]
135            points = np.array(shape_annotation["points"], dtype=float)
136            rows, columns = draw_polygon(points[:, 1], points[:, 0], shape=shape)
137            labels[rows, columns] = label_id
138
139        imageio.imwrite(label_path, labels, compression="zlib")
140
141
142def get_cystods_data(path: Union[os.PathLike, str], download: bool = False) -> str:
143    """Download the CystoDS dataset.
144
145    Args:
146        path: Filepath to a folder where the downloaded data will be saved.
147        download: Whether to download the data if it is not present.
148
149    Returns:
150        The filepath to the folder where the data is stored.
151    """
152    os.makedirs(path, exist_ok=True)
153
154    csv_path = os.path.join(path, "cystods.csv")
155    util.download_source(path=csv_path, url=CSV_URL, download=download, checksum=CSV_CHECKSUM)
156
157    filenames = _get_segmented_filenames(csv_path)
158
159    image_dir = os.path.join(path, "images")
160    segmentation_dir = os.path.join(path, "segmentations")
161    have_images = os.path.exists(image_dir) and len(glob(os.path.join(image_dir, "*.png"))) == len(filenames)
162    have_segmentations = os.path.exists(segmentation_dir) and \
163        len(glob(os.path.join(segmentation_dir, "*.json"))) == len(filenames)
164    if have_images and have_segmentations:
165        return path
166
167    if not download:
168        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
169
170    os.makedirs(image_dir, exist_ok=True)
171    os.makedirs(segmentation_dir, exist_ok=True)
172
173    image_urls = _get_osf_folder_download_urls(IMAGES_FOLDER_ID)
174    segmentation_urls = _get_osf_folder_download_urls(SEGMENTATIONS_FOLDER_ID)
175
176    for fname in tqdm(filenames, desc="Download the CystoDS images and segmentations"):
177        name = os.path.splitext(fname)[0]
178        segmentation_name = f"{name}.json"
179
180        if fname not in image_urls or segmentation_name not in segmentation_urls:
181            warn(f"'{fname}' is listed as segmented in 'cystods.csv' but missing from the OSF folders.")
182            continue
183
184        image_path = os.path.join(image_dir, fname)
185        _download_with_retries(image_path, image_urls[fname], download)
186
187        segmentation_path = os.path.join(segmentation_dir, segmentation_name)
188        _download_with_retries(segmentation_path, segmentation_urls[segmentation_name], download)
189
190    return path
191
192
193def get_cystods_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
194    """Get paths to the CystoDS data.
195
196    Args:
197        path: Filepath to a folder where the downloaded data will be saved.
198        download: Whether to download the data if it is not present.
199
200    Returns:
201        List of filepaths for the image data.
202        List of filepaths for the label data.
203    """
204    data_dir = get_cystods_data(path, download)
205
206    filenames = _get_segmented_filenames(os.path.join(data_dir, "cystods.csv"))
207    # A few filenames may be missing an image or a segmentation on the OSF side (see 'get_cystods_data'),
208    # so only keep the ones that were actually downloaded.
209    filenames = [
210        fname for fname in filenames
211        if os.path.exists(os.path.join(data_dir, "images", fname))
212        and os.path.exists(os.path.join(data_dir, "segmentations", f"{os.path.splitext(fname)[0]}.json"))
213    ]
214
215    label_dir = os.path.join(data_dir, "labels")
216    _rasterize_labels(os.path.join(data_dir, "segmentations"), filenames, label_dir)
217
218    raw_paths = natsorted(os.path.join(data_dir, "images", fname) for fname in filenames)
219    label_paths = natsorted(glob(os.path.join(label_dir, "*.tif")))
220
221    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
222    assert all(
223        os.path.splitext(os.path.basename(raw_path))[0] == os.path.splitext(os.path.basename(label_path))[0]
224        for raw_path, label_path in zip(raw_paths, label_paths)
225    )
226
227    return raw_paths, label_paths
228
229
230def get_cystods_dataset(
231    path: Union[os.PathLike, str],
232    patch_shape: Tuple[int, int],
233    resize_inputs: bool = False,
234    download: bool = False,
235    **kwargs,
236) -> Dataset:
237    """Get the CystoDS dataset for bladder lesion and landmark segmentation.
238
239    Args:
240        path: Filepath to a folder where the downloaded data will be saved.
241        patch_shape: The patch shape to use for training.
242        resize_inputs: Whether to resize the inputs.
243        download: Whether to download the data if it is not present.
244        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
245
246    Returns:
247        The segmentation dataset.
248    """
249    raw_paths, label_paths = get_cystods_paths(path, download)
250
251    if resize_inputs:
252        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
253        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
254            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
255        )
256
257    return torch_em.default_segmentation_dataset(
258        raw_paths=raw_paths,
259        raw_key=None,
260        label_paths=label_paths,
261        label_key=None,
262        patch_shape=patch_shape,
263        is_seg_dataset=False,
264        **kwargs,
265    )
266
267
268def get_cystods_loader(
269    path: Union[os.PathLike, str],
270    batch_size: int,
271    patch_shape: Tuple[int, int],
272    resize_inputs: bool = False,
273    download: bool = False,
274    **kwargs,
275) -> DataLoader:
276    """Get the CystoDS dataloader for bladder lesion and landmark segmentation.
277
278    Args:
279        path: Filepath to a folder where the downloaded data will be saved.
280        batch_size: The batch size for training.
281        patch_shape: The patch shape to use for training.
282        resize_inputs: Whether to resize the inputs.
283        download: Whether to download the data if it is not present.
284        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
285
286    Returns:
287        The DataLoader.
288    """
289    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
290    dataset = get_cystods_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
291    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)