torch_em.data.datasets.medical.fuseg

  1"""The FUSeg dataset contains annotations for wound segmentation in clinical foot ulcer photographs.
  2
  3The dataset consists of 1,210 foot ulcer images collected over two years from 889 patients for the
  42021 MICCAI Foot Ulcer Segmentation (FUSeg) Challenge. Only the 'train' and 'validation' splits are
  5used here, as they are the only splits with publicly released ground-truth masks; the 'test' split
  6ground truth is kept private by the challenge organizers for the leaderboard.
  7
  8The dataset is located at https://github.com/uwm-bigdata/wound-segmentation.
  9The dataset is from the publication https://doi.org/10.3390/info15030140.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14import subprocess
 15from glob import glob
 16from natsort import natsorted
 17from typing import Union, Tuple, Literal, List
 18
 19import imageio.v3 as imageio
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26from ..light_microscopy.neurips_cell_seg import to_rgb
 27
 28
 29URL = "https://github.com/uwm-bigdata/wound-segmentation.git"
 30
 31CHALLENGE_DIR = "Foot Ulcer Segmentation Challenge"
 32
 33
 34def get_fuseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the FUSeg dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the data is downloaded for further processing.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        Filepath where the data is downloaded.
 43    """
 44    data_dir = os.path.join(path, "wound-segmentation", "data", CHALLENGE_DIR)
 45    if os.path.exists(data_dir):
 46        return data_dir
 47
 48    if not download:
 49        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
 50
 51    os.makedirs(path, exist_ok=True)
 52    repo_dir = os.path.join(path, "wound-segmentation")
 53    subprocess.run(
 54        ["git", "clone", "--filter=blob:none", "--sparse", "--depth", "1", "--quiet", URL, repo_dir], check=True
 55    )
 56    subprocess.run(["git", "sparse-checkout", "set", os.path.join("data", CHALLENGE_DIR)], cwd=repo_dir, check=True)
 57
 58    return data_dir
 59
 60
 61def _binarize_mask(mask_path, out_path):
 62    if os.path.exists(out_path):
 63        return
 64    mask = imageio.imread(mask_path)
 65    if mask.ndim == 3:
 66        mask = mask[..., 0]
 67    mask = (mask > 127).astype("uint8")
 68    imageio.imwrite(out_path, mask)
 69
 70
 71def get_fuseg_paths(
 72    path: Union[os.PathLike, str], split: Literal["train", "validation"] = "train", download: bool = False
 73) -> Tuple[List[str], List[str]]:
 74    """Get paths to the FUSeg data.
 75
 76    Args:
 77        path: Filepath to a folder where the data is downloaded for further processing.
 78        split: The choice of data split. Either 'train' or 'validation'.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        List of filepaths for the image data.
 83        List of filepaths for the label data.
 84    """
 85    if split not in ["train", "validation"]:
 86        raise ValueError(f"'{split}' is not a valid split. Choose 'train' or 'validation'.")
 87
 88    data_dir = get_fuseg_data(path, download)
 89
 90    image_dir = os.path.join(data_dir, split, "images")
 91    label_dir = os.path.join(data_dir, split, "labels")
 92    preprocessed_dir = os.path.join(data_dir, split, "preprocessed_labels")
 93    os.makedirs(preprocessed_dir, exist_ok=True)
 94
 95    image_paths = natsorted(glob(os.path.join(image_dir, "*.png")))
 96
 97    label_paths = []
 98    for image_path in image_paths:
 99        fname = os.path.basename(image_path)
100        mask_path = os.path.join(label_dir, fname)
101        out_path = os.path.join(preprocessed_dir, fname)
102        _binarize_mask(mask_path, out_path)
103        label_paths.append(out_path)
104
105    return image_paths, label_paths
106
107
108def get_fuseg_dataset(
109    path: Union[os.PathLike, str],
110    patch_shape: Tuple[int, int],
111    split: Literal["train", "validation"] = "train",
112    resize_inputs: bool = False,
113    download: bool = False,
114    **kwargs
115) -> Dataset:
116    """Get the FUSeg dataset for wound segmentation in foot ulcer photographs.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        patch_shape: The patch shape to use for training.
121        split: The choice of data split. Either 'train' or 'validation'.
122        resize_inputs: Whether to resize the inputs to the expected patch shape.
123        download: Whether to download the data if it is not present.
124        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
125
126    Returns:
127        The segmentation dataset.
128    """
129    image_paths, label_paths = get_fuseg_paths(path, split, download)
130
131    if resize_inputs:
132        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
133        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
134            kwargs=kwargs,
135            patch_shape=patch_shape,
136            resize_inputs=resize_inputs,
137            resize_kwargs=resize_kwargs,
138            ensure_rgb=to_rgb,
139        )
140
141    return torch_em.default_segmentation_dataset(
142        raw_paths=image_paths,
143        raw_key=None,
144        label_paths=label_paths,
145        label_key=None,
146        patch_shape=patch_shape,
147        is_seg_dataset=False,
148        **kwargs
149    )
150
151
152def get_fuseg_loader(
153    path: Union[os.PathLike, str],
154    batch_size: int,
155    patch_shape: Tuple[int, int],
156    split: Literal["train", "validation"] = "train",
157    resize_inputs: bool = False,
158    download: bool = False,
159    **kwargs
160) -> DataLoader:
161    """Get the FUSeg dataloader for wound segmentation in foot ulcer photographs.
162
163    Args:
164        path: Filepath to a folder where the data is downloaded for further processing.
165        batch_size: The batch size for training.
166        patch_shape: The patch shape to use for training.
167        split: The choice of data split. Either 'train' or 'validation'.
168        resize_inputs: Whether to resize the inputs to the expected patch shape.
169        download: Whether to download the data if it is not present.
170        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
171
172    Returns:
173        The DataLoader.
174    """
175    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
176    dataset = get_fuseg_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
177    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)