torch_em.data.datasets.medical.woives

  1"""The WOIVES dataset contains annotations for retinal vessel segmentation in ultra-widefield
  2swept-source optical coherence tomography angiography (SS-OCTA) images.
  3
  4The dataset consists of 206 eyes from 152 participants. Each eye has a superficial-retina OCTA image
  5(1536x1280 pixels, 24x20 mm field of view) paired with a probabilistic soft-label vessel annotation. The
  6soft labels are the pixel-wise average of five binary annotations that were made with SAM-OCTA assistance
  7and then refined by two senior ophthalmologists. They are stored as uint8 images in [0, 255], where
  8value / 255 is the vessel probability. The release ships no binary masks, and it also contains
  9deep/full-retina and choroid slabs, fundus photographs and SLO images, which this loader does not use.
 10
 11By default the soft labels are thresholded into binary vessel masks (see the 'binarize_threshold'
 12argument). Please report the threshold you use. With 'binarize_threshold=None' the soft probability
 13maps in [0, 1] are returned as float labels instead.
 14
 15The official split is a subject-level five-fold cross-validation, selected with the 'fold' and 'split'
 16arguments.
 17
 18The data is located at https://doi.org/10.5281/zenodo.21904672, released under a CC-BY-4.0 license.
 19
 20This dataset is from the publication https://arxiv.org/abs/2609.12574.
 21Please cite it if you use this dataset for your research.
 22"""
 23
 24import os
 25import json
 26from glob import glob
 27from tqdm import tqdm
 28from natsort import natsorted
 29from typing import Union, Tuple, Literal, List, Optional
 30
 31import numpy as np
 32
 33from torch.utils.data import Dataset, DataLoader
 34
 35import torch_em
 36
 37from .. import util
 38
 39
 40URL = "https://zenodo.org/api/records/21904672/files/WOIVES_v1.0.zip/content"
 41CHECKSUM = "e921284a2b2dc5348607439ead2b567456f73881a5c79251176d700b91f6c757"
 42
 43SPLITS = ["train", "val", "test"]
 44FOLDS = [0, 1, 2, 3, 4]
 45
 46
 47def _convert_labels(label_dir, out_dir, binarize_threshold):
 48    import tifffile
 49    from PIL import Image
 50
 51    os.makedirs(out_dir, exist_ok=True)
 52    for label_path in tqdm(natsorted(glob(os.path.join(label_dir, "mask_*.png"))), desc="Convert the WOIVES labels"):
 53        out_path = os.path.join(out_dir, os.path.basename(label_path)[len("mask_"):].replace(".png", ".tif"))
 54        if os.path.exists(out_path):
 55            continue
 56
 57        soft = np.asarray(Image.open(label_path), dtype="float32") / 255.0
 58        label = soft if binarize_threshold is None else (soft >= binarize_threshold).astype("uint8")
 59
 60        tmp_path = f"{out_path}.{os.getpid()}.incomplete.tif"
 61        tifffile.imwrite(tmp_path, label, compression="zlib")
 62        os.replace(tmp_path, out_path)
 63
 64
 65def get_woives_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 66    """Download the WOIVES dataset.
 67
 68    Args:
 69        path: Filepath to a folder where the data is downloaded for further processing.
 70        download: Whether to download the data if it is not present.
 71
 72    Returns:
 73        Filepath to the extracted dataset folder.
 74    """
 75    data_dir = os.path.join(path, "WOIVES_v1.0")
 76    if os.path.exists(data_dir):
 77        return data_dir
 78
 79    os.makedirs(path, exist_ok=True)
 80
 81    zip_path = os.path.join(path, "WOIVES_v1.0.zip")
 82    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 83    util.unzip(zip_path=zip_path, dst=path, remove=False)
 84
 85    assert os.path.exists(data_dir), f"The extraction of the WOIVES archive did not create '{data_dir}'."
 86
 87    return data_dir
 88
 89
 90def get_woives_paths(
 91    path: Union[os.PathLike, str],
 92    split: Literal["train", "val", "test"],
 93    fold: int = 0,
 94    binarize_threshold: Optional[float] = 0.5,
 95    download: bool = False,
 96) -> Tuple[List[str], List[str]]:
 97    """Get paths to the WOIVES data.
 98
 99    Args:
100        path: Filepath to a folder where the data is downloaded for further processing.
101        split: The choice of data split. One of 'train', 'val' or 'test'.
102        fold: The fold of the official five-fold cross-validation. One of 0 to 4.
103        binarize_threshold: The threshold on the vessel probability for creating binary masks (vessel if the
104            probability is at least the threshold). Set to None to use the soft probability maps as labels.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        List of filepaths for the image data.
109        List of filepaths for the label data.
110    """
111    if split not in SPLITS:
112        raise ValueError(f"'{split}' is not a valid split. Choose one of {SPLITS}.")
113    if fold not in FOLDS:
114        raise ValueError(f"'{fold}' is not a valid fold. Choose one of {FOLDS}.")
115    if binarize_threshold is not None and not 0.0 < binarize_threshold <= 1.0:
116        raise ValueError(f"The threshold must be in (0, 1] or None, but got {binarize_threshold}.")
117
118    data_dir = get_woives_data(path, download)
119
120    label_name = "labels_soft" if binarize_threshold is None else f"labels_bin_{binarize_threshold:g}"
121    label_dir = os.path.join(path, label_name)
122    _convert_labels(os.path.join(data_dir, "Label"), label_dir, binarize_threshold)
123
124    with open(os.path.join(data_dir, "splits", "fold_split.json")) as f:
125        names = json.load(f)["folds"][f"fold{fold}"][split]
126
127    raw_paths = [os.path.join(data_dir, "Image", "OCTA", "OCTA_superficial retina", name) for name in names]
128    label_paths = [os.path.join(label_dir, name.replace(".png", ".tif")) for name in names]
129
130    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
131    assert all(os.path.exists(p) for p in raw_paths + label_paths)
132
133    return raw_paths, label_paths
134
135
136def get_woives_dataset(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int],
139    split: Literal["train", "val", "test"],
140    fold: int = 0,
141    binarize_threshold: Optional[float] = 0.5,
142    resize_inputs: bool = False,
143    download: bool = False,
144    **kwargs
145) -> Dataset:
146    """Get the WOIVES dataset for retinal vessel segmentation in ultra-widefield OCTA images.
147
148    Args:
149        path: Filepath to a folder where the data is downloaded for further processing.
150        patch_shape: The patch shape to use for training.
151        split: The choice of data split. One of 'train', 'val' or 'test'.
152        fold: The fold of the official five-fold cross-validation. One of 0 to 4.
153        binarize_threshold: The threshold on the vessel probability for creating binary masks.
154            Set to None to use the soft probability maps as labels.
155        resize_inputs: Whether to resize the inputs to the patch shape.
156        download: Whether to download the data if it is not present.
157        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
158
159    Returns:
160        The segmentation dataset.
161    """
162    raw_paths, label_paths = get_woives_paths(path, split, fold, binarize_threshold, download)
163
164    if resize_inputs:
165        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
166        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
167            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
168        )
169
170    return torch_em.default_segmentation_dataset(
171        raw_paths=raw_paths,
172        raw_key=None,
173        label_paths=label_paths,
174        label_key=None,
175        is_seg_dataset=False,
176        patch_shape=patch_shape,
177        **kwargs
178    )
179
180
181def get_woives_loader(
182    path: Union[os.PathLike, str],
183    batch_size: int,
184    patch_shape: Tuple[int, int],
185    split: Literal["train", "val", "test"],
186    fold: int = 0,
187    binarize_threshold: Optional[float] = 0.5,
188    resize_inputs: bool = False,
189    download: bool = False,
190    **kwargs
191) -> DataLoader:
192    """Get the WOIVES dataloader for retinal vessel segmentation in ultra-widefield OCTA images.
193
194    Args:
195        path: Filepath to a folder where the data is downloaded for further processing.
196        batch_size: The batch size for training.
197        patch_shape: The patch shape to use for training.
198        split: The choice of data split. One of 'train', 'val' or 'test'.
199        fold: The fold of the official five-fold cross-validation. One of 0 to 4.
200        binarize_threshold: The threshold on the vessel probability for creating binary masks.
201            Set to None to use the soft probability maps as labels.
202        resize_inputs: Whether to resize the inputs to the patch shape.
203        download: Whether to download the data if it is not present.
204        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
205
206    Returns:
207        The DataLoader.
208    """
209    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
210    dataset = get_woives_dataset(
211        path, patch_shape, split, fold, binarize_threshold, resize_inputs, download, **ds_kwargs
212    )
213    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)