torch_em.data.datasets.medical.periorbital_seg

  1"""The Periorbital Segmentation dataset contains annotations for periorbital anatomical
  2structures (eyebrow, sclera, iris, caruncle, eyelid) in cropped eye images.
  3
  4The dataset bundles both the cropped eye images and the segmentation masks (no need to
  5source images separately from the Chicago Face Database or CelebAMask-HQ). It combines two
  6subsets:
  7- 'celeb': 2015 image-mask pairs cropped from CelebAMask-HQ.
  8- 'cfd': 827 image-mask pairs cropped from the Chicago Face Database.
  9(The archive also has a 'combined_final_data' folder, but it is inconsistent - it has 3670
 10images and only 2842 masks - so this module sources images and masks from the per-origin
 11'celeb_final_data' and 'cfd_final_data' folders instead, which are fully paired.)
 12
 13The label legend (confirmed from the dataset's own 'coco2voc_aux.py' preprocessing script) is:
 14- background: 0, eyebrow: 1, sclera: 2, iris (incl. pupil): 3, caruncle: 4, eyelid: 5
 15
 16The dataset is located at https://doi.org/10.5281/zenodo.13916845 (Zenodo, CC BY 4.0).
 17This dataset is from the publication https://doi.org/10.48550/arXiv.2409.20407.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22from glob import glob
 23from natsort import natsorted
 24from typing import Union, Tuple, Literal, List
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL = "https://zenodo.org/records/13916845/files/periorbital_dataset.zip?download=1"
 34CHECKSUM = "7c1cd92e4bc58e5b288c6eeb1769a98fa119338bfb25a2d3bb7ccebed8d99011"
 35
 36SUBSETS = ["celeb", "cfd"]
 37
 38
 39def get_periorbital_seg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 40    """Download the Periorbital Segmentation dataset.
 41
 42    Args:
 43        path: Filepath to a folder where the data is downloaded for further processing.
 44        download: Whether to download the data if it is not present.
 45
 46    Returns:
 47        Filepath where the data is stored.
 48    """
 49    data_dir = os.path.join(path, "periorbital_dataset")
 50    if os.path.exists(data_dir):
 51        return data_dir
 52
 53    os.makedirs(path, exist_ok=True)
 54
 55    zip_path = os.path.join(path, "periorbital_dataset.zip")
 56    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 57    util.unzip(zip_path=zip_path, dst=path)
 58
 59    return data_dir
 60
 61
 62def get_periorbital_seg_paths(
 63    path: Union[os.PathLike, str],
 64    subset: Union[Literal["celeb", "cfd"], List[str]] = SUBSETS,
 65    download: bool = False,
 66) -> Tuple[List[str], List[str]]:
 67    """Get paths to the Periorbital Segmentation data.
 68
 69    Args:
 70        path: Filepath to a folder where the data is downloaded for further processing.
 71        subset: The choice of data subset(s). Either 'celeb', 'cfd', or a list of both. By default, loads both.
 72        download: Whether to download the data if it is not present.
 73
 74    Returns:
 75        List of filepaths for the image data.
 76        List of filepaths for the label data.
 77    """
 78    data_dir = get_periorbital_seg_data(path, download)
 79
 80    subsets = [subset] if isinstance(subset, str) else subset
 81    for s in subsets:
 82        if s not in SUBSETS:
 83            raise ValueError(f"'{s}' is not a valid subset. Please choose from {SUBSETS}.")
 84
 85    image_paths, gt_paths = [], []
 86    for s in subsets:
 87        this_image_paths = natsorted(glob(os.path.join(data_dir, f"{s}_final_data", f"{s}_output_images", "*.jpg")))
 88        this_gt_dir = os.path.join(data_dir, f"{s}_final_data", f"{s}_output_masks")
 89        for image_path in this_image_paths:
 90            fname = os.path.splitext(os.path.basename(image_path))[0]
 91            gt_path = os.path.join(this_gt_dir, f"{fname}.png")
 92            if not os.path.exists(gt_path):
 93                continue
 94            image_paths.append(image_path)
 95            gt_paths.append(gt_path)
 96
 97    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0, (
 98        "No image-mask pairs were found. The expected per-subset 'celeb_output_images' / 'celeb_output_masks' "
 99        f"(and 'cfd_output_images' / 'cfd_output_masks') layout may not match the actual structure of the "
100        f"downloaded data. Please inspect the data at '{data_dir}' and update the search pattern accordingly."
101    )
102
103    return image_paths, gt_paths
104
105
106def get_periorbital_seg_dataset(
107    path: Union[os.PathLike, str],
108    patch_shape: Tuple[int, int],
109    subset: Union[Literal["celeb", "cfd"], List[str]] = SUBSETS,
110    resize_inputs: bool = False,
111    download: bool = False,
112    **kwargs
113) -> Dataset:
114    """Get the Periorbital Segmentation dataset for periorbital anatomical structure segmentation.
115
116    Args:
117        path: Filepath to a folder where the data is downloaded for further processing.
118        patch_shape: The patch shape to use for training.
119        subset: The choice of data subset(s). Either 'celeb', 'cfd', or a list of both. By default, loads both.
120        resize_inputs: Whether to resize inputs to the desired patch shape.
121        download: Whether to download the data if it is not present.
122        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
123
124    Returns:
125        The segmentation dataset.
126    """
127    image_paths, gt_paths = get_periorbital_seg_paths(path, subset, download)
128
129    if resize_inputs:
130        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
131        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
132            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
133        )
134
135    return torch_em.default_segmentation_dataset(
136        raw_paths=image_paths,
137        raw_key=None,
138        label_paths=gt_paths,
139        label_key=None,
140        is_seg_dataset=False,
141        patch_shape=patch_shape,
142        **kwargs
143    )
144
145
146def get_periorbital_seg_loader(
147    path: Union[os.PathLike, str],
148    batch_size: int,
149    patch_shape: Tuple[int, int],
150    subset: Union[Literal["celeb", "cfd"], List[str]] = SUBSETS,
151    resize_inputs: bool = False,
152    download: bool = False,
153    **kwargs
154) -> DataLoader:
155    """Get the Periorbital Segmentation dataloader for periorbital anatomical structure segmentation.
156
157    Args:
158        path: Filepath to a folder where the data is downloaded for further processing.
159        batch_size: The batch size for training.
160        patch_shape: The patch shape to use for training.
161        subset: The choice of data subset(s). Either 'celeb', 'cfd', or a list of both. By default, loads both.
162        resize_inputs: Whether to resize inputs to the desired patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
165
166    Returns:
167        The DataLoader.
168    """
169    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
170    dataset = get_periorbital_seg_dataset(path, patch_shape, subset, resize_inputs, download, **ds_kwargs)
171    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)