torch_em.data.datasets.medical.ocutox

  1"""The OCUTOX dataset contains annotations for active and inactive lesion segmentation
  2in fundus images of patients with ocular toxoplasmosis.
  3
  4The dataset contains 412 fundus images collected at the Hospital de Clinicas and the
  5Hospital General Pediatrico Acosta Nu medical centers in Asuncion, Paraguay, of which
  6280 images (with active and / or inactive toxoplasmosis lesions) have pixel-level lesion
  7masks delineated by ophthalmologists. The remaining images are labeled 'healthy' and do
  8not have a lesion mask.
  9
 10This dataset is located at https://doi.org/10.5281/zenodo.5156940 (CC BY 4.0).
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15import re
 16from glob import glob
 17from typing import Union, Tuple, List
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26URL = "https://zenodo.org/records/5156940/files/Ocular_Toxoplasmosis_Data_V3.zip"
 27CHECKSUM = "d088dcfe6c678923eed20e18052935cff70a3e0ede8f87f06406caa734727ece"
 28
 29# Masks carry suffixes for lesion sub-regions of the same image, eg. '-a' (active lesion),
 30# '-i' (inactive lesion) and numeric variants ('-2', '-3', '-a-2', ...).
 31MASK_SUFFIX_PATTERN = re.compile(r"-(?:a|i)(?:-\d+)?$|-\d+$", flags=re.IGNORECASE)
 32
 33
 34def get_ocutox_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the OCUTOX 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, "images")
 45    if os.path.exists(data_dir):
 46        return path
 47
 48    os.makedirs(path, exist_ok=True)
 49
 50    zip_path = os.path.join(path, "Ocular_Toxoplasmosis_Data_V3.zip")
 51    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 52    util.unzip(zip_path=zip_path, dst=path)
 53
 54    return path
 55
 56
 57def get_ocutox_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 58    """Get paths to the OCUTOX data.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        download: Whether to download the data if it is not present.
 63
 64    Returns:
 65        List of filepaths for the image data.
 66        List of filepaths for the label data.
 67    """
 68    data_dir = get_ocutox_data(path=path, download=download)
 69
 70    image_dir = os.path.join(data_dir, "images")
 71    mask_dir = os.path.join(data_dir, "masks")
 72
 73    image_files = {os.path.basename(p).lower(): p for p in glob(os.path.join(image_dir, "*.*"))}
 74    mask_paths = sorted(glob(os.path.join(mask_dir, "*.*")))
 75
 76    image_paths, matched_mask_paths = [], []
 77    for mask_path in mask_paths:
 78        fname = os.path.basename(mask_path)
 79        stem, ext = os.path.splitext(fname.lower())
 80        base_stem = MASK_SUFFIX_PATTERN.sub("", stem)
 81        image_path = image_files.get(base_stem + ext)
 82        if image_path is None:
 83            raise RuntimeError(f"Could not find the matching image for the mask at '{mask_path}'.")
 84
 85        image_paths.append(image_path)
 86        matched_mask_paths.append(mask_path)
 87
 88    if len(image_paths) == 0 or len(image_paths) != len(matched_mask_paths):
 89        raise RuntimeError("Something went wrong with fetching the image and label paths.")
 90
 91    return image_paths, matched_mask_paths
 92
 93
 94def get_ocutox_dataset(
 95    path: Union[os.PathLike, str],
 96    patch_shape: Tuple[int, int],
 97    resize_inputs: bool = False,
 98    download: bool = False,
 99    **kwargs
100) -> Dataset:
101    """Get the OCUTOX dataset for segmentation of ocular toxoplasmosis lesions in fundus images.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        patch_shape: The patch shape to use for training.
106        resize_inputs: Whether to resize the inputs to the expected patch shape.
107        download: Whether to download the data if it is not present.
108        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
109
110    Returns:
111        The segmentation dataset.
112    """
113    image_paths, gt_paths = get_ocutox_paths(path, download)
114
115    if resize_inputs:
116        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
117        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
118            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
119        )
120
121    return torch_em.default_segmentation_dataset(
122        raw_paths=image_paths,
123        raw_key=None,
124        label_paths=gt_paths,
125        label_key=None,
126        patch_shape=patch_shape,
127        is_seg_dataset=False,
128        **kwargs
129    )
130
131
132def get_ocutox_loader(
133    path: Union[os.PathLike, str],
134    batch_size: int,
135    patch_shape: Tuple[int, int],
136    resize_inputs: bool = False,
137    download: bool = False,
138    **kwargs
139) -> DataLoader:
140    """Get the OCUTOX dataloader for segmentation of ocular toxoplasmosis lesions in fundus images.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        batch_size: The batch size for training.
145        patch_shape: The patch shape to use for training.
146        resize_inputs: Whether to resize the inputs to the expected patch shape.
147        download: Whether to download the data if it is not present.
148        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
149
150    Returns:
151        The DataLoader.
152    """
153    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
154    dataset = get_ocutox_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
155    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)