torch_em.data.datasets.medical.amd_dme_3d_oct

  1"""The AMD-DME-3D-OCT dataset contains annotations for pigment epithelial detachment (PED) and
  2intraretinal fluid (IRF) segmentation in full 3D optical coherence tomography (OCT) volumes.
  3
  4The dataset consists of 224 volumes (122 AMD and 102 DME), of which 104 (62 AMD, 42 DME) come with
  5volumetric annotations. The remaining 120 volumes are unlabeled and are not used by this loader.
  6Each annotated volume was independently segmented by three junior ophthalmologists, and the
  7resulting masks were reviewed and corrected by an experienced ophthalmologist. The label volumes
  8are binary (PED and IRF are merged into a single foreground class), matching each 512x512x512 raw
  9volume ('<case><eye><picture>.tif') to its mask ('<case><eye><picture>_mask.tif').
 10
 11NOTE: This is a distinct, volumetric (3D) dataset. It is not the same as the already-integrated
 12`torch_em.data.datasets.medical.amd_sd`, which ships 2D B-scan annotations for wet AMD lesions only.
 13
 14The data is located at https://doi.org/10.6084/m9.figshare.30582035, released under a CC-BY-4.0 license.
 15NOTE: The archive is a single ~19.6 GB 7z file, which requires the 'p7zip' CLI to extract
 16(see `torch_em.data.datasets.util.unzip_7z`).
 17
 18This dataset is from the publication https://doi.org/10.1038/s41597-025-06497-1.
 19Please cite it if you use this dataset for your research.
 20"""
 21
 22import os
 23from glob import glob
 24from natsort import natsorted
 25from typing import Union, Tuple, Literal, List
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URL = "https://ndownloader.figshare.com/files/59443823"
 35CHECKSUM = "31d1e2c581be375654a3d757eaa2710f4e91628b0779c39f2a5d04afe4ba0318"
 36
 37DISEASES = ["AMD", "DME"]
 38
 39
 40def get_amd_dme_3d_oct_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 41    """Download the AMD-DME-3D-OCT dataset.
 42
 43    Args:
 44        path: Filepath to a folder where the data is downloaded for further processing.
 45        download: Whether to download the data if it is not present.
 46
 47    Returns:
 48        Filepath where the data is downloaded.
 49    """
 50    data_dir = os.path.join(path, "AMD_DME_3D_Dataset")
 51    if all(os.path.exists(os.path.join(data_dir, f"{disease}_labeled")) for disease in DISEASES):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55
 56    archive_path = os.path.join(path, "AMD_DME_3D_Dataset.7z")
 57    util.download_source(path=archive_path, url=URL, download=download, checksum=CHECKSUM)
 58    util.unzip_7z(path_7z=archive_path, dst=path, remove=False)
 59
 60    assert all(os.path.exists(os.path.join(data_dir, f"{disease}_labeled")) for disease in DISEASES), \
 61        f"The extraction of the AMD-DME-3D-OCT archive did not create the expected folders in '{data_dir}'."
 62
 63    return data_dir
 64
 65
 66def get_amd_dme_3d_oct_paths(
 67    path: Union[os.PathLike, str], disease: Literal["AMD", "DME"] = "AMD", download: bool = False,
 68) -> Tuple[List[str], List[str]]:
 69    """Get paths to the AMD-DME-3D-OCT data.
 70
 71    Args:
 72        path: Filepath to a folder where the data is downloaded for further processing.
 73        disease: The choice of disease. Either 'AMD' or 'DME'.
 74        download: Whether to download the data if it is not present.
 75
 76    Returns:
 77        List of filepaths for the image data.
 78        List of filepaths for the label data.
 79    """
 80    if disease not in DISEASES:
 81        raise ValueError(f"'{disease}' is not a valid choice of disease. Choose one of {DISEASES}.")
 82
 83    data_dir = get_amd_dme_3d_oct_data(path, download)
 84
 85    label_paths = natsorted(glob(os.path.join(data_dir, f"{disease}_labeled", "label", "*_mask.tif")))
 86    raw_paths = [
 87        os.path.join(data_dir, f"{disease}_labeled", "input", os.path.basename(p).replace("_mask.tif", ".tif"))
 88        for p in label_paths
 89    ]
 90
 91    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 92    assert all(os.path.exists(p) for p in raw_paths)
 93
 94    return raw_paths, label_paths
 95
 96
 97def get_amd_dme_3d_oct_dataset(
 98    path: Union[os.PathLike, str],
 99    patch_shape: Tuple[int, int, int],
100    disease: Literal["AMD", "DME"] = "AMD",
101    resize_inputs: bool = False,
102    download: bool = False,
103    **kwargs
104) -> Dataset:
105    """Get the AMD-DME-3D-OCT dataset for PED and IRF segmentation in 3D OCT volumes.
106
107    Args:
108        path: Filepath to a folder where the data is downloaded for further processing.
109        patch_shape: The patch shape to use for training.
110        disease: The choice of disease. Either 'AMD' or 'DME'.
111        resize_inputs: Whether to resize the inputs to the patch shape.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
114
115    Returns:
116        The segmentation dataset.
117    """
118    raw_paths, label_paths = get_amd_dme_3d_oct_paths(path, disease, download)
119
120    if resize_inputs:
121        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
122        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
123            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
124        )
125
126    return torch_em.default_segmentation_dataset(
127        raw_paths=raw_paths,
128        raw_key=None,
129        label_paths=label_paths,
130        label_key=None,
131        is_seg_dataset=True,
132        patch_shape=patch_shape,
133        ndim=3,
134        **kwargs
135    )
136
137
138def get_amd_dme_3d_oct_loader(
139    path: Union[os.PathLike, str],
140    batch_size: int,
141    patch_shape: Tuple[int, int, int],
142    disease: Literal["AMD", "DME"] = "AMD",
143    resize_inputs: bool = False,
144    download: bool = False,
145    **kwargs
146) -> DataLoader:
147    """Get the AMD-DME-3D-OCT dataloader for PED and IRF segmentation in 3D OCT volumes.
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        disease: The choice of disease. Either 'AMD' or 'DME'.
154        resize_inputs: Whether to resize the inputs to the patch shape.
155        download: Whether to download the data if it is not present.
156        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
157
158    Returns:
159        The DataLoader.
160    """
161    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
162    dataset = get_amd_dme_3d_oct_dataset(path, patch_shape, disease, resize_inputs, download, **ds_kwargs)
163    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)