torch_em.data.datasets.medical.longitudinal_mh_oct

  1"""This dataset contains pixel-level segmented longitudinal OCT scans of idiopathic
  2full-thickness macular hole (iFTMH) surgery outcomes, collected at one preoperative
  3('baseline') and six postoperative time points ('2weeks', '3months', '6months',
  4'12months', '24months', '48months').
  5
  6The dataset contains 2,591 fovea-centered horizontal and vertical B-scans from 493
  7patients, with expert-validated pixel-level segmentation masks for 12 retinal
  8structures and pathologies, both anatomical (e.g. external limiting membrane,
  9ellipsoid zone, retinal pigment epithelium) and pathological (e.g. macular hole,
 10cysts, epiretinal membrane, subretinal fluid).
 11
 12NOTE: This is distinct from the OIMHS dataset (see 'torch_em/data/datasets/medical/
 13oimhs.py'), which is a different, non-longitudinal macular hole OCT dataset.
 14
 15The dataset is located at https://doi.org/10.6084/m9.figshare.32605218 and is
 16licensed under CC-BY-4.0.
 17
 18This dataset is from the publication https://doi.org/10.1038/s41597-026-08154-7.
 19Please cite it if you use this dataset for your research.
 20"""
 21
 22import os
 23from glob import glob
 24from pathlib import Path
 25from natsort import natsorted
 26from typing import Union, Tuple, List, Optional, Literal
 27
 28import imageio.v3 as imageio
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URL = "https://ndownloader.figshare.com/files/65374524"
 38CHECKSUM = "e9956c1587123e94342cc74dd9d2a50c9c833ffe52457688862ece8b849a4c97"
 39
 40TIMEPOINTS = ["baseline", "2weeks", "3months", "6months", "12months", "24months", "48months"]
 41
 42
 43def get_longitudinal_mh_oct_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 44    """Download the longitudinal macular hole OCT data.
 45
 46    Args:
 47        path: Filepath to a folder where the data is downloaded for further processing.
 48        download: Whether to download the data if it is not present.
 49
 50    Returns:
 51        Filepath where the data is downloaded.
 52    """
 53    data_dir = os.path.join(path, "Dataset")
 54    if os.path.exists(data_dir):
 55        return data_dir
 56
 57    os.makedirs(path, exist_ok=True)
 58
 59    zip_path = os.path.join(path, "Dataset.zip")
 60    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 61    util.unzip(zip_path=zip_path, dst=path)
 62
 63    return data_dir
 64
 65
 66def get_longitudinal_mh_oct_paths(
 67    path: Union[os.PathLike, str],
 68    timepoint: Optional[Literal["baseline", "2weeks", "3months", "6months", "12months", "24months", "48months"]] = None,  # noqa
 69    download: bool = False,
 70) -> Tuple[List[str], List[str]]:
 71    """Get paths to the longitudinal macular hole OCT data.
 72
 73    Args:
 74        path: Filepath to a folder where the data is downloaded for further processing.
 75        timepoint: The choice of a specific postoperative (or 'baseline') timepoint. By default, loads all
 76            timepoints.
 77        download: Whether to download the data if it is not present.
 78
 79    Returns:
 80        List of filepaths for the image data.
 81        List of filepaths for the label data.
 82    """
 83    data_dir = get_longitudinal_mh_oct_data(path, download)
 84
 85    timepoints = TIMEPOINTS if timepoint is None else [timepoint]
 86
 87    pp_dir = os.path.join(data_dir, "preprocessed_images")
 88    os.makedirs(pp_dir, exist_ok=True)
 89
 90    image_paths, gt_paths = [], []
 91    for this_timepoint in timepoints:
 92        this_gt_paths = natsorted(
 93            glob(os.path.join(data_dir, this_timepoint, f"{this_timepoint}_Masks", "*.png"))
 94        )
 95        for gt_path in this_gt_paths:
 96            org_image_path = os.path.join(
 97                data_dir, this_timepoint, f"{this_timepoint}_OCT", f"{Path(gt_path).stem}.tiff"
 98            )
 99            if not os.path.exists(org_image_path):
100                continue
101
102            # The raw B-scans are RGBA tiffs, but 'ImageCollectionDataset' expects RGB inputs.
103            # The alpha channel is dropped once here and the result cached as a '.tif' file.
104            image_path = os.path.join(pp_dir, f"{this_timepoint}_{Path(gt_path).stem}.tif")
105            if not os.path.exists(image_path):
106                image = imageio.imread(org_image_path)
107                if image.ndim == 3 and image.shape[-1] == 4:
108                    image = image[..., :3]
109                imageio.imwrite(image_path, image, compression="zlib")
110
111            image_paths.append(image_path)
112            gt_paths.append(gt_path)
113
114    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0, (
115        "No image-mask pairs were found. The expected per-timepoint '<timepoint>_OCT' / '<timepoint>_Masks' "
116        "folder layout may not match the actual structure of the downloaded data. Please inspect the data at "
117        f"'{data_dir}'."
118    )
119
120    return image_paths, gt_paths
121
122
123def get_longitudinal_mh_oct_dataset(
124    path: Union[os.PathLike, str],
125    patch_shape: Tuple[int, int],
126    timepoint: Optional[Literal["baseline", "2weeks", "3months", "6months", "12months", "24months", "48months"]] = None,  # noqa
127    resize_inputs: bool = False,
128    download: bool = False,
129    **kwargs
130) -> Dataset:
131    """Get the longitudinal macular hole OCT dataset for segmentation of 12 retinal structures and pathologies.
132
133    Args:
134        path: Filepath to a folder where the data is downloaded for further processing.
135        patch_shape: The patch shape to use for training.
136        timepoint: The choice of a specific postoperative (or 'baseline') timepoint. By default, loads all
137            timepoints.
138        resize_inputs: Whether to resize inputs to the desired patch shape.
139        download: Whether to download the data if it is not present.
140        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
141
142    Returns:
143        The segmentation dataset.
144    """
145    image_paths, gt_paths = get_longitudinal_mh_oct_paths(path, timepoint, download)
146
147    if resize_inputs:
148        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
149        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
150            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
151        )
152
153    return torch_em.default_segmentation_dataset(
154        raw_paths=image_paths,
155        raw_key=None,
156        label_paths=gt_paths,
157        label_key=None,
158        patch_shape=patch_shape,
159        is_seg_dataset=False,
160        **kwargs
161    )
162
163
164def get_longitudinal_mh_oct_loader(
165    path: Union[os.PathLike, str],
166    batch_size: int,
167    patch_shape: Tuple[int, int],
168    timepoint: Optional[Literal["baseline", "2weeks", "3months", "6months", "12months", "24months", "48months"]] = None,  # noqa
169    resize_inputs: bool = False,
170    download: bool = False,
171    **kwargs
172) -> DataLoader:
173    """Get the longitudinal macular hole OCT dataloader for segmentation of 12 retinal structures and pathologies.
174
175    Args:
176        path: Filepath to a folder where the data is downloaded for further processing.
177        batch_size: The batch size for training.
178        patch_shape: The patch shape to use for training.
179        timepoint: The choice of a specific postoperative (or 'baseline') timepoint. By default, loads all
180            timepoints.
181        resize_inputs: Whether to resize inputs to the desired patch shape.
182        download: Whether to download the data if it is not present.
183        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
184
185    Returns:
186        The DataLoader.
187    """
188    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
189    dataset = get_longitudinal_mh_oct_dataset(path, patch_shape, timepoint, resize_inputs, download, **ds_kwargs)
190    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)