torch_em.data.datasets.medical.intraretinal_cystoid_fluid

  1"""The Intraretinal Cystoid Fluid dataset contains annotations for cystoid macular edema (CME)
  2segmentation in optical coherence tomography (OCT) images of diabetic macular edema (DME) patients.
  3
  4The dataset is located at https://www.kaggle.com/datasets/zeeshanahmed13/intraretinal-cystoid-fluid.
  5This dataset is from the publication https://doi.org/10.1002/ima.22662.
  6Please cite it if you use this dataset for your research.
  7"""
  8
  9import os
 10import re
 11from glob import glob
 12from tqdm import tqdm
 13from pathlib import Path
 14from typing import Union, Tuple, List
 15
 16import imageio.v3 as imageio
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25def get_intraretinal_cystoid_fluid_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 26    """Download the Intraretinal Cystoid Fluid dataset.
 27
 28    Args:
 29        path: Filepath to a folder where the data is downloaded for further processing.
 30        download: Whether to download the data if it is not present.
 31
 32    Returns:
 33        Filepath where the data is downloaded.
 34    """
 35    data_dir = os.path.join(path, "2021-training-data-ZA", "2021-training-data-ZA")
 36    if os.path.exists(data_dir):
 37        return data_dir
 38
 39    os.makedirs(path, exist_ok=True)
 40
 41    util.download_source_kaggle(
 42        path=path, dataset_name="zeeshanahmed13/intraretinal-cystoid-fluid", download=download,
 43    )
 44    zip_path = os.path.join(path, "intraretinal-cystoid-fluid.zip")
 45    util.unzip(zip_path=zip_path, dst=path)
 46
 47    return data_dir
 48
 49
 50def get_intraretinal_cystoid_fluid_paths(
 51    path: Union[os.PathLike, str], download: bool = False
 52) -> Tuple[List[str], List[str]]:
 53    """Get paths to the Intraretinal Cystoid Fluid data.
 54
 55    Args:
 56        path: Filepath to a folder where the data is downloaded for further processing.
 57        download: Whether to download the data if it is not present.
 58
 59    Returns:
 60        List of filepaths for the image data.
 61        List of filepaths for the label data.
 62    """
 63    data_dir = get_intraretinal_cystoid_fluid_data(path=path, download=download)
 64
 65    neu_gt_dir = os.path.join(path, "preprocessed_masks")
 66    os.makedirs(neu_gt_dir, exist_ok=True)
 67
 68    sample_dirs = sorted(glob(os.path.join(data_dir, "*")))
 69
 70    image_paths, gt_paths = [], []
 71    for sample_dir in tqdm(sample_dirs, desc="Preprocessing labels"):
 72        if not os.path.isdir(sample_dir):
 73            continue
 74
 75        image_path = glob(os.path.join(sample_dir, "images", "*"))[0]
 76        mask_paths = sorted(glob(os.path.join(sample_dir, "masks", "*")))
 77
 78        # A single sample folder ('815DME_F') ships a stray extra mask file that belongs to another
 79        # sample ('MASK-DME511.png'). Filter to the mask matching the sample's leading numeric id.
 80        if len(mask_paths) > 1:
 81            sample_name = os.path.basename(sample_dir)
 82            prefix_match = re.match(r"^(\d+)", sample_name)
 83            if prefix_match is not None:
 84                prefix = prefix_match.group(1)
 85                mask_paths = [p for p in mask_paths if os.path.basename(p).startswith(prefix)]
 86
 87        mask_path = mask_paths[0]
 88
 89        neu_gt_path = os.path.join(neu_gt_dir, f"{Path(mask_path).stem}.tif")
 90        if not os.path.exists(neu_gt_path):
 91            gt = imageio.imread(mask_path)
 92            gt = (gt > 0).astype("uint8")
 93            imageio.imwrite(neu_gt_path, gt, compression="zlib")
 94
 95        image_paths.append(image_path)
 96        gt_paths.append(neu_gt_path)
 97
 98    return image_paths, gt_paths
 99
100
101def get_intraretinal_cystoid_fluid_dataset(
102    path: Union[os.PathLike, str],
103    patch_shape: Tuple[int, int],
104    resize_inputs: bool = False,
105    download: bool = False,
106    **kwargs
107) -> Dataset:
108    """Get the Intraretinal Cystoid Fluid dataset for CME segmentation in OCT images.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        patch_shape: The patch shape to use for training.
113        resize_inputs: Whether to resize the inputs to the patch shape.
114        download: Whether to download the data if it is not present.
115        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
116
117    Returns:
118        The segmentation dataset.
119    """
120    image_paths, gt_paths = get_intraretinal_cystoid_fluid_paths(path, download)
121
122    if resize_inputs:
123        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
124        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
125            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
126        )
127
128    return torch_em.default_segmentation_dataset(
129        raw_paths=image_paths,
130        raw_key=None,
131        label_paths=gt_paths,
132        label_key=None,
133        patch_shape=patch_shape,
134        is_seg_dataset=False,
135        **kwargs
136    )
137
138
139def get_intraretinal_cystoid_fluid_loader(
140    path: Union[os.PathLike, str],
141    batch_size: int,
142    patch_shape: Tuple[int, int],
143    resize_inputs: bool = False,
144    download: bool = False,
145    **kwargs
146) -> DataLoader:
147    """Get the Intraretinal Cystoid Fluid dataloader for CME segmentation in OCT images.
148
149    Args:
150        path: Filepath to a folder where the data is downloaded for further processing.
151        batch_size: The batch size for training.
152        patch_shape: The patch shape to use for training.
153        resize_inputs: Whether to resize the inputs to the patch shape.
154        download: Whether to download the data if it is not present.
155        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
156
157    Returns:
158        The DataLoader.
159    """
160    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
161    dataset = get_intraretinal_cystoid_fluid_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
162    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)