torch_em.data.datasets.medical.chaksu

  1"""The Chaksu dataset contains annotations for optic disc and optic cup
  2segmentation in Fundus images.
  3
  4This dataset is located at https://doi.org/10.6084/m9.figshare.20123135.v2, under the CC BY 4.0 license.
  5The dataset is from the publication https://doi.org/10.1038/s41597-023-01943-4.
  6Please cite it if you use this dataset for your research.
  7
  8NOTE: The full archive also ships four other expert annotation and fusion variants at full
  9uncompressed resolution, expanding to over 200GB. Only the raw fundus images and the
 10STAPLE-fused consensus masks are downloaded here, as those are the ones needed for training.
 11"""
 12
 13import io
 14import os
 15import zipfile
 16from glob import glob
 17from typing import Union, Tuple, Literal, List
 18
 19import requests
 20from tqdm import tqdm
 21import imageio.v3 as imageio
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URLS = {
 31    "train": "https://ndownloader.figshare.com/files/37875672",
 32    "test": "https://ndownloader.figshare.com/files/37875687",
 33}
 34DEVICES = ["Bosch", "Forus", "Remidio"]
 35
 36
 37class _RemoteZipFile(io.RawIOBase):
 38    """File-like wrapper that reads a remote zip archive via HTTP range requests."""
 39
 40    def __init__(self, url):
 41        self.url = url
 42        self.session = requests.Session()
 43        response = self.session.get(url, headers={"Range": "bytes=0-0"}, timeout=30)
 44        self.size = int(response.headers["Content-Range"].split("/")[-1])
 45        self.pos = 0
 46
 47    def readable(self):
 48        return True
 49
 50    def seekable(self):
 51        return True
 52
 53    def seek(self, offset, whence=0):
 54        if whence == 0:
 55            self.pos = offset
 56        elif whence == 1:
 57            self.pos += offset
 58        elif whence == 2:
 59            self.pos = self.size + offset
 60        return self.pos
 61
 62    def tell(self):
 63        return self.pos
 64
 65    def readinto(self, buffer):
 66        end = min(self.pos + len(buffer), self.size) - 1
 67        if end < self.pos:
 68            return 0
 69        response = self.session.get(self.url, headers={"Range": f"bytes={self.pos}-{end}"}, timeout=60)
 70        data = response.content
 71        buffer[:len(data)] = data
 72        self.pos += len(data)
 73        return len(data)
 74
 75
 76def _download_split(url, dst, split):
 77    remote_file = io.BufferedReader(_RemoteZipFile(url), buffer_size=1 << 20)
 78    with zipfile.ZipFile(remote_file) as zf:
 79        members = [
 80            name for name in zf.namelist() if not name.endswith("/") and "__MACOSX" not in name
 81            and not name.endswith(".DS_Store")
 82            and ("1.0_Original_Fundus_Images/" in name
 83                 or ("5.0_OD_OC_Mean_Median_Majority_STAPLE/" in name and "/STAPLE/" in name))
 84        ]
 85        for member in tqdm(members, desc=f"Downloading Chaksu '{split}' split"):
 86            zf.extract(member, dst)
 87
 88
 89def get_chaksu_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 90    """Download the Chaksu dataset.
 91
 92    Args:
 93        path: Filepath to a folder where the data is downloaded for further processing.
 94        download: Whether to download the data if it is not present.
 95
 96    Returns:
 97        Filepath where the data is downloaded.
 98    """
 99    os.makedirs(path, exist_ok=True)
100    for split, url in URLS.items():
101        split_dir = os.path.join(path, split.capitalize())
102        if os.path.exists(split_dir):
103            continue
104        if not download:
105            raise RuntimeError(f"Chaksu data is not found at {split_dir} and 'download' is set to False.")
106        _download_split(url, path, split)
107
108    return path
109
110
111def _binarize_mask(gt_path, gt_dir):
112    dst_path = os.path.join(gt_dir, os.path.basename(gt_path))
113    if os.path.exists(dst_path):
114        return dst_path
115
116    os.makedirs(gt_dir, exist_ok=True)
117    mask = imageio.imread(gt_path)[..., 0] > 128  # the STAPLE masks are near-binary grayscale stored as RGBA
118    imageio.imwrite(dst_path, mask.astype("uint8"))
119    return dst_path
120
121
122def get_chaksu_paths(
123    path: Union[os.PathLike, str],
124    split: Literal['train', 'test'],
125    task: Literal["optic_disc", "optic_cup"] = "optic_disc",
126    download: bool = False,
127) -> Tuple[List[str], List[str]]:
128    """Get paths to the Chaksu data.
129
130    Args:
131        path: Filepath to a folder where the data is downloaded for further processing.
132        split: The choice of data split.
133        task: The choice of labels for the specific task.
134        download: Whether to download the data if it is not present.
135
136    Returns:
137        List of filepaths for the image data.
138        List of filepaths for the label data.
139    """
140    data_dir = get_chaksu_data(path=path, download=download)
141
142    assert split in ["train", "test"], f"'{split}' is not a valid split."
143    assert task in ["optic_disc", "optic_cup"], f"'{task}' is not a valid task."
144
145    split_dir = split.capitalize()
146    region = "Disc" if task == "optic_disc" else "Cup"
147
148    image_paths, gt_paths = [], []
149    for device in DEVICES:
150        device_image_paths = sorted(glob(os.path.join(data_dir, split_dir, "1.0_Original_Fundus_Images", device, "*")))
151        gt_dir = os.path.join(data_dir, split_dir, "segmentation_masks", device, region)
152        for image_path in device_image_paths:
153            stem = os.path.splitext(os.path.basename(image_path))[0]
154            staple_path = os.path.join(
155                data_dir, split_dir, "5.0_OD_OC_Mean_Median_Majority_STAPLE", device, region, "STAPLE", f"{stem}.png"
156            )
157            if not os.path.exists(staple_path):
158                continue
159
160            image_paths.append(image_path)
161            gt_paths.append(_binarize_mask(staple_path, gt_dir))
162
163    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0
164
165    return image_paths, gt_paths
166
167
168def get_chaksu_dataset(
169    path: Union[os.PathLike, str],
170    patch_shape: Tuple[int, int],
171    split: Literal['train', 'test'],
172    task: Literal["optic_disc", "optic_cup"] = "optic_disc",
173    resize_inputs: bool = False,
174    download: bool = False,
175    **kwargs
176) -> Dataset:
177    """Get the Chaksu dataset for segmentation of optic disc and optic cup in fundus images.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        patch_shape: The patch shape to use for training.
182        split: The choice of data split.
183        task: The choice of labels for the specific task.
184        resize_inputs: Whether to resize the inputs to the expected patch shape.
185        download: Whether to download the data if it is not present.
186        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
187
188    Returns:
189        The segmentation dataset.
190    """
191    image_paths, gt_paths = get_chaksu_paths(path, split, task, download)
192
193    if resize_inputs:
194        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
195        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
196            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
197        )
198
199    return torch_em.default_segmentation_dataset(
200        raw_paths=image_paths,
201        raw_key=None,
202        label_paths=gt_paths,
203        label_key=None,
204        patch_shape=patch_shape,
205        is_seg_dataset=False,
206        **kwargs
207    )
208
209
210def get_chaksu_loader(
211    path: Union[os.PathLike, str],
212    batch_size: int,
213    patch_shape: Tuple[int, int],
214    split: Literal['train', 'test'],
215    task: Literal["optic_disc", "optic_cup"] = "optic_disc",
216    resize_inputs: bool = False,
217    download: bool = False,
218    **kwargs
219) -> DataLoader:
220    """Get the Chaksu dataloader for segmentation of optic disc and optic cup in fundus images.
221
222    Args:
223        path: Filepath to a folder where the data is downloaded for further processing.
224        batch_size: The batch size for training.
225        patch_shape: The patch shape to use for training.
226        split: The choice of data split.
227        task: The choice of labels for the specific task.
228        resize_inputs: Whether to resize the inputs to the expected patch shape.
229        download: Whether to download the data if it is not present.
230        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
231
232    Returns:
233        The DataLoader.
234    """
235    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
236    dataset = get_chaksu_dataset(path, patch_shape, split, task, resize_inputs, download, **ds_kwargs)
237    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)