torch_em.data.datasets.medical.deeprt

  1"""The DeepRT dataset contains annotations for retinal tissue segmentation in optical coherence
  2tomography (OCT) B-scans, originally released to support self-supervised pretraining for diabetic
  3retinopathy classification.
  4
  5The dataset ships 1,286 OCT images (acquired with Spectralis and Topcon devices) with binary retinal
  6tissue masks, pre-split into 'train', 'validation' and 'test' (744 / 265 / 277 images respectively)
  7via the 'file_names_complete' mapping files shipped in the archive.
  8NOTE: the associated publication reports 1,009 "semantic segmentations", a smaller curated subset;
  9this loader exposes the full 1,286 labeled images shipped in the 'thickness_segmentation_data' archive.
 10
 11The data is located at https://doi.org/10.5281/zenodo.3626020, released under a CC-BY-4.0 license.
 12
 13This dataset is from the publication https://doi.org/10.1038/s42256-020-00247-1.
 14Please cite it if you use this dataset for your research.
 15"""
 16
 17import os
 18from glob import glob
 19from typing import Union, Tuple, List, Literal
 20
 21import numpy as np
 22import pandas as pd
 23from PIL import Image
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URL = "https://zenodo.org/records/3626020/files/thickness_segmentation_data.tar.gz"
 33CHECKSUM = "f0c3c0ce9f140e470f6579bc6ec634d719d5947726f8b38f48f04f88026cdaa8"
 34
 35SPLITS = ["train", "validation", "test"]
 36
 37
 38def _binarize_labels(data_dir):
 39    label_paths = glob(os.path.join(data_dir, "data", "all_labels", "*.png"))
 40    for label_path in label_paths:
 41        label = np.array(Image.open(label_path).convert("L"))
 42        Image.fromarray((label > 0).astype("uint8")).save(label_path)
 43
 44
 45def get_deeprt_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 46    """Download the DeepRT dataset.
 47
 48    Args:
 49        path: Filepath to a folder where the data is downloaded for further processing.
 50        download: Whether to download the data if it is not present.
 51
 52    Returns:
 53        Filepath where the data is downloaded.
 54    """
 55    data_dir = os.path.join(path, "DeepRT")
 56    if os.path.exists(data_dir):
 57        return data_dir
 58
 59    os.makedirs(path, exist_ok=True)
 60
 61    tar_path = os.path.join(path, "thickness_segmentation_data.tar.gz")
 62    util.download_source(path=tar_path, url=URL, download=download, checksum=CHECKSUM)
 63    util.unzip_tarfile(tar_path=tar_path, dst=data_dir, remove=False)
 64
 65    _binarize_labels(data_dir)
 66
 67    return data_dir
 68
 69
 70def get_deeprt_paths(
 71    path: Union[os.PathLike, str], split: Literal["train", "validation", "test"] = "train", download: bool = False,
 72) -> Tuple[List[str], List[str]]:
 73    """Get paths to the DeepRT data.
 74
 75    Args:
 76        path: Filepath to a folder where the data is downloaded for further processing.
 77        split: The choice of data split. Either 'train', 'validation' or 'test'.
 78        download: Whether to download the data if it is not present.
 79
 80    Returns:
 81        List of filepaths for the image data.
 82        List of filepaths for the label data.
 83    """
 84    if split not in SPLITS:
 85        raise ValueError(f"'{split}' is not a valid split. Choose one of {SPLITS}.")
 86
 87    data_dir = get_deeprt_data(path, download)
 88
 89    mapping = pd.read_csv(os.path.join(data_dir, "data", "file_names_complete", f"{split}_new_old_mapping.csv"))
 90    ids = mapping["new_id"].tolist()
 91
 92    raw_paths = [os.path.join(data_dir, "data", "all_images", f"{i}.png") for i in ids]
 93    label_paths = [os.path.join(data_dir, "data", "all_labels", f"{i}.png") for i in ids]
 94
 95    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 96    assert all(os.path.exists(p) for p in raw_paths) and all(os.path.exists(p) for p in label_paths)
 97
 98    return raw_paths, label_paths
 99
100
101def get_deeprt_dataset(
102    path: Union[os.PathLike, str],
103    patch_shape: Tuple[int, int],
104    split: Literal["train", "validation", "test"] = "train",
105    resize_inputs: bool = False,
106    download: bool = False,
107    **kwargs
108) -> Dataset:
109    """Get the DeepRT dataset for retinal tissue segmentation in OCT images.
110
111    Args:
112        path: Filepath to a folder where the data is downloaded for further processing.
113        patch_shape: The patch shape to use for training.
114        split: The choice of data split. Either 'train', 'validation' or 'test'.
115        resize_inputs: Whether to resize the inputs to the patch shape.
116        download: Whether to download the data if it is not present.
117        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
118
119    Returns:
120        The segmentation dataset.
121    """
122    raw_paths, label_paths = get_deeprt_paths(path, split, download)
123
124    if resize_inputs:
125        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
126        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
127            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
128        )
129
130    return torch_em.default_segmentation_dataset(
131        raw_paths=raw_paths,
132        raw_key=None,
133        label_paths=label_paths,
134        label_key=None,
135        is_seg_dataset=False,
136        patch_shape=patch_shape,
137        ndim=2,
138        **kwargs
139    )
140
141
142def get_deeprt_loader(
143    path: Union[os.PathLike, str],
144    batch_size: int,
145    patch_shape: Tuple[int, int],
146    split: Literal["train", "validation", "test"] = "train",
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> DataLoader:
151    """Get the DeepRT dataloader for retinal tissue segmentation in OCT images.
152
153    Args:
154        path: Filepath to a folder where the data is downloaded for further processing.
155        batch_size: The batch size for training.
156        patch_shape: The patch shape to use for training.
157        split: The choice of data split. Either 'train', 'validation' or 'test'.
158        resize_inputs: Whether to resize the inputs to the patch shape.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
161
162    Returns:
163        The DataLoader.
164    """
165    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
166    dataset = get_deeprt_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
167    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)