torch_em.data.datasets.medical.tear_meniscus

  1"""The Tear Meniscus dataset contains annotations for pixel-level segmentation of the lower tear meniscus
  2in external-eye colour and infrared images.
  3
  4The dataset comprises 1693 colour images and 1739 infrared images, collected from five clinical centres
  5across China. The initial annotations were produced with a human-computer-interactive approach and were
  6subsequently reviewed and corrected by a senior ophthalmologist, so the ground truth is clinician-reviewed.
  7The shipped label maps are binary (0 = background, 1 = tear meniscus); some of them are stored as
  83-channel (RGB) images with all channels holding the same grayscale value, which this module reduces to
  9a single-channel binary mask. The raw images also mix formats (RGBA, RGB and grayscale across the
 10different centres), which this module normalizes to 3-channel RGB.
 11
 12The dataset is located at https://doi.org/10.6084/m9.figshare.28650536.v2 (CC BY 4.0).
 13
 14This dataset is from the publication https://doi.org/10.1038/s41597-025-06460-0.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from tqdm import tqdm
 21from pathlib import Path
 22from natsort import natsorted
 23from typing import Union, Tuple, Literal, List
 24
 25import numpy as np
 26import imageio.v3 as imageio
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://ndownloader.figshare.com/files/56406179"
 36CHECKSUM = "b4614214e69098b09160f4713c5e0f11fcd6a7a158a5217ca9cd15ba2687f170"
 37
 38LABEL_IDS = {"background": 0, "tear_meniscus": 1}
 39
 40MODALITIES = {"colour": "Colour", "infrared": "Infrared"}
 41
 42
 43def get_tear_meniscus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 44    """Download the Tear Meniscus dataset.
 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, "Open DataSet2")
 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, "Open 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 _preprocess_images(raw_paths: List[str]) -> List[str]:
 67    """Normalize the (RGBA / RGB / grayscale) raw images to 3-channel RGB."""
 68    preprocessed_paths = []
 69    for raw_path in tqdm(raw_paths, desc="Preprocessing images"):
 70        parent = Path(raw_path).parent.parent
 71        preprocessed_dir = os.path.join(parent, "Original_rgb")
 72        os.makedirs(preprocessed_dir, exist_ok=True)
 73
 74        preprocessed_path = os.path.join(preprocessed_dir, f"{Path(raw_path).stem}.tif")
 75        preprocessed_paths.append(preprocessed_path)
 76        if os.path.exists(preprocessed_path):
 77            continue
 78
 79        raw = imageio.imread(raw_path)
 80        if raw.ndim == 2:
 81            raw = np.repeat(raw[..., None], 3, axis=-1)
 82        else:
 83            raw = raw[..., :3]
 84        imageio.imwrite(preprocessed_path, raw, compression="zlib")
 85
 86    return preprocessed_paths
 87
 88
 89def _preprocess_labels(raw_label_paths: List[str]) -> List[str]:
 90    """Reduce the (partially 3-channel) label images to single-channel binary masks."""
 91    preprocessed_paths = []
 92    for label_path in tqdm(raw_label_paths, desc="Preprocessing labels"):
 93        parent = Path(label_path).parent.parent
 94        preprocessed_dir = os.path.join(parent, "Label_binary")
 95        os.makedirs(preprocessed_dir, exist_ok=True)
 96
 97        preprocessed_path = os.path.join(preprocessed_dir, f"{Path(label_path).stem}.tif")
 98        preprocessed_paths.append(preprocessed_path)
 99        if os.path.exists(preprocessed_path):
100            continue
101
102        label = imageio.imread(label_path)
103        if label.ndim == 3:
104            label = label[..., 0]
105        label = (label > 127).astype("uint8")
106        imageio.imwrite(preprocessed_path, label, compression="zlib")
107
108    return preprocessed_paths
109
110
111def get_tear_meniscus_paths(
112    path: Union[os.PathLike, str],
113    modality: Literal["colour", "infrared"] = "colour",
114    download: bool = False,
115) -> Tuple[List[str], List[str]]:
116    """Get paths to the Tear Meniscus data.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        modality: The imaging modality. Either 'colour' or 'infrared'.
121        download: Whether to download the data if it is not present.
122
123    Returns:
124        List of filepaths for the image data.
125        List of filepaths for the label data.
126    """
127    data_dir = get_tear_meniscus_data(path, download)
128
129    if modality not in MODALITIES:
130        raise ValueError(f"'{modality}' is not a valid modality. Please choose one of {list(MODALITIES)}.")
131    prefix = MODALITIES[modality]
132
133    raw_image_paths = natsorted(glob(os.path.join(data_dir, f"{prefix}*", "Original", "*.PNG")))
134    raw_label_paths = [p.replace(f"{os.sep}Original{os.sep}", f"{os.sep}Label{os.sep}") for p in raw_image_paths]
135
136    assert len(raw_image_paths) > 0, f"Could not find any images in '{data_dir}'."
137    for label_path in raw_label_paths:
138        assert os.path.exists(label_path), label_path
139
140    image_paths = _preprocess_images(raw_image_paths)
141    label_paths = _preprocess_labels(raw_label_paths)
142
143    return image_paths, label_paths
144
145
146def get_tear_meniscus_dataset(
147    path: Union[os.PathLike, str],
148    patch_shape: Tuple[int, int],
149    modality: Literal["colour", "infrared"] = "colour",
150    resize_inputs: bool = False,
151    download: bool = False,
152    **kwargs
153) -> Dataset:
154    """Get the Tear Meniscus dataset for segmentation of the tear meniscus.
155
156    Args:
157        path: Filepath to a folder where the data is downloaded for further processing.
158        patch_shape: The patch shape to use for training.
159        modality: The imaging modality. Either 'colour' or 'infrared'.
160        resize_inputs: Whether to resize the inputs to the expected patch shape.
161        download: Whether to download the data if it is not present.
162        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
163
164    Returns:
165        The segmentation dataset.
166    """
167    image_paths, label_paths = get_tear_meniscus_paths(path, modality, download)
168
169    if resize_inputs:
170        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
171        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
172            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
173        )
174
175    return torch_em.default_segmentation_dataset(
176        raw_paths=image_paths,
177        raw_key=None,
178        label_paths=label_paths,
179        label_key=None,
180        patch_shape=patch_shape,
181        is_seg_dataset=False,
182        **kwargs
183    )
184
185
186def get_tear_meniscus_loader(
187    path: Union[os.PathLike, str],
188    batch_size: int,
189    patch_shape: Tuple[int, int],
190    modality: Literal["colour", "infrared"] = "colour",
191    resize_inputs: bool = False,
192    download: bool = False,
193    **kwargs
194) -> DataLoader:
195    """Get the Tear Meniscus dataloader for segmentation of the tear meniscus.
196
197    Args:
198        path: Filepath to a folder where the data is downloaded for further processing.
199        batch_size: The batch size for training.
200        patch_shape: The patch shape to use for training.
201        modality: The imaging modality. Either 'colour' or 'infrared'.
202        resize_inputs: Whether to resize the inputs to the expected patch shape.
203        download: Whether to download the data if it is not present.
204        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
205
206    Returns:
207        The DataLoader.
208    """
209    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
210    dataset = get_tear_meniscus_dataset(path, patch_shape, modality, resize_inputs, download, **ds_kwargs)
211    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)