torch_em.data.datasets.medical.cadis

  1"""The CaDIS dataset contains annotations for anatomy and surgical instrument segmentation
  2in RGB video frames of cataract surgery.
  3
  4The dataset consists of 4670 frames sampled from 25 videos of the training set of the CATARACTS
  5challenge (https://cataracts.grand-challenge.org). Each frame is densely labeled at the pixel level
  6into up to 36 semantic classes: 4 anatomical structures (pupil, iris, cornea, skin), 29 surgical
  7instrument classes and 3 other object classes (surgical tape, hand, eye retractors). The labels are
  8stored as single-channel images, where the pixel value is the class index.
  9
 10NOTE: The official train/val/test split assigns videos 5, 7 and 16 to validation and videos 2, 12
 11and 22 to test, with the remaining videos used for training (see `VAL_VIDEOS` and `TEST_VIDEOS`).
 12
 13NOTE: The dataset is hosted on the CATARACTS Grand Challenge platform at
 14https://cataracts.grand-challenge.org/CaDIS/ and requires registration, so it cannot be downloaded
 15automatically. As of writing this, the direct download link on that page is not yet public
 16("Uploading - Link coming up soon"), so the data must currently be requested from the challenge
 17organizers. To obtain it:
 18- Visit https://cataracts.grand-challenge.org/CaDIS/ and register / log in via the challenge portal.
 19- Follow the instructions on the 'Data' / 'Download' pages of the challenge to request access to CaDIS.
 20- Once you have the data, place it such that each video folder is located at
 21  '<path>/CADIS/segmentation/Video<NN>/Images/*' and '<path>/CADIS/segmentation/Video<NN>/Labels/*',
 22  matching the structure of the original release.
 23
 24This dataset is from the publication https://doi.org/10.1016/j.media.2021.102053.
 25Please cite it if you use this dataset in your research.
 26"""
 27
 28import os
 29import re
 30from glob import glob
 31from natsort import natsorted
 32from typing import Union, Tuple, List, Literal
 33
 34from torch.utils.data import Dataset, DataLoader
 35
 36import torch_em
 37
 38from .. import util
 39
 40
 41VAL_VIDEOS = [5, 7, 16]
 42TEST_VIDEOS = [2, 12, 22]
 43
 44IMAGE_EXTENSIONS = ("*.png", "*.jpg", "*.jpeg")
 45
 46
 47def _video_number(video_dir):
 48    match = re.search(r"(\d+)", os.path.basename(video_dir))
 49    if match is None:
 50        raise ValueError(f"Could not parse a video number from '{video_dir}'.")
 51    return int(match.group(1))
 52
 53
 54def get_cadis_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 55    """Obtain the CaDIS dataset.
 56
 57    Args:
 58        path: Filepath to a folder where the manually downloaded data is stored.
 59        download: Whether to download the data if it is not present. The data cannot be downloaded
 60            automatically, so this raises if the data has not been downloaded manually.
 61
 62    Returns:
 63        Filepath to the folder with the per-video 'Images' and 'Labels' subfolders.
 64    """
 65    data_dir = os.path.join(path, "CADIS", "segmentation")
 66    if os.path.exists(data_dir) and glob(os.path.join(data_dir, "Video*")):
 67        return data_dir
 68
 69    msg = "'torch_em' cannot download this dataset, because CaDIS is distributed via the CATARACTS Grand "
 70    msg += "Challenge and requires registration. Please visit 'https://cataracts.grand-challenge.org/CaDIS/', "
 71    msg += "register / log in and request access to the data via the challenge's 'Data' page, then place it "
 72    msg += f"such that video folders are located at '{data_dir}/Video<NN>/Images' and '.../Video<NN>/Labels'."
 73    raise NotImplementedError(msg)
 74
 75
 76def get_cadis_paths(
 77    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
 78) -> Tuple[List[str], List[str]]:
 79    """Get paths to the CaDIS data.
 80
 81    Args:
 82        path: Filepath to a folder where the manually downloaded data is stored.
 83        split: The choice of data split.
 84        download: Whether to download the data if it is not present.
 85
 86    Returns:
 87        List of filepaths for the image data.
 88        List of filepaths for the label data.
 89    """
 90    data_dir = get_cadis_data(path, download)
 91
 92    video_dirs = natsorted(glob(os.path.join(data_dir, "Video*")))
 93    assert video_dirs, f"Did not find any 'Video*' folders in '{data_dir}'."
 94
 95    if split == "val":
 96        video_dirs = [v for v in video_dirs if _video_number(v) in VAL_VIDEOS]
 97    elif split == "test":
 98        video_dirs = [v for v in video_dirs if _video_number(v) in TEST_VIDEOS]
 99    elif split == "train":
100        video_dirs = [v for v in video_dirs if _video_number(v) not in VAL_VIDEOS + TEST_VIDEOS]
101    else:
102        raise ValueError(f"'{split}' is not a valid split.")
103
104    image_paths, label_paths = [], []
105    for video_dir in video_dirs:
106        vimage_paths = []
107        for ext in IMAGE_EXTENSIONS:
108            vimage_paths.extend(glob(os.path.join(video_dir, "Images", ext)))
109        vimage_paths = natsorted(vimage_paths)
110
111        vlabel_paths = natsorted(glob(os.path.join(video_dir, "Labels", "*.png")))
112
113        assert vimage_paths and len(vimage_paths) == len(vlabel_paths), \
114            f"The images and labels for '{video_dir}' do not match."
115
116        image_paths.extend(vimage_paths)
117        label_paths.extend(vlabel_paths)
118
119    return image_paths, label_paths
120
121
122def get_cadis_dataset(
123    path: Union[os.PathLike, str],
124    patch_shape: Tuple[int, int],
125    split: Literal["train", "val", "test"],
126    resize_inputs: bool = False,
127    download: bool = False,
128    **kwargs
129) -> Dataset:
130    """Get the CaDIS dataset for anatomy and surgical instrument segmentation.
131
132    Args:
133        path: Filepath to a folder where the manually downloaded data is stored.
134        patch_shape: The patch shape to use for training.
135        split: The choice of data split.
136        resize_inputs: Whether to resize inputs to the desired patch shape.
137        download: Whether to download the data if it is not present.
138        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
139
140    Returns:
141        The segmentation dataset.
142    """
143    image_paths, label_paths = get_cadis_paths(path, split, download)
144
145    if resize_inputs:
146        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
147        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
148            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
149        )
150
151    return torch_em.default_segmentation_dataset(
152        raw_paths=image_paths,
153        raw_key=None,
154        label_paths=label_paths,
155        label_key=None,
156        is_seg_dataset=False,
157        patch_shape=patch_shape,
158        **kwargs
159    )
160
161
162def get_cadis_loader(
163    path: Union[os.PathLike, str],
164    batch_size: int,
165    patch_shape: Tuple[int, int],
166    split: Literal["train", "val", "test"],
167    resize_inputs: bool = False,
168    download: bool = False,
169    **kwargs
170) -> DataLoader:
171    """Get the CaDIS dataloader for anatomy and surgical instrument segmentation.
172
173    Args:
174        path: Filepath to a folder where the manually downloaded data is stored.
175        batch_size: The batch size for training.
176        patch_shape: The patch shape to use for training.
177        split: The choice of data split.
178        resize_inputs: Whether to resize inputs to the desired patch shape.
179        download: Whether to download the data if it is not present.
180        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
181
182    Returns:
183        The DataLoader.
184    """
185    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
186    dataset = get_cadis_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
187    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)