torch_em.data.datasets.medical.oct5k

  1"""OCT5k contains pixel-wise multi-grader annotations of retinal layer bands in optical coherence
  2tomography (OCT) B-scans, covering AMD (age-related macular degeneration), DME (diabetic macular
  3edema) and healthy subjects.
  4
  5The masks are located at https://doi.org/10.5522/04/22128671 (UCL Research Data Repository, a
  6figshare instance), released under a CC0 license. This dataset is from the publication
  7https://doi.org/10.1038/s41597-024-04259-z. Please cite it if you use this dataset in your research.
  8
  9NOTE: The raw OCT images are not shipped with the masks archive. The dataset's own scripts fetch them
 10from a third-party source instead (R. Rasti's Macular OCT dataset), distributed as a password-protected
 11archive on Google Drive (password published on the page linked from the OCT5k scripts, no registration
 12required). As of this writing, the primary Google Drive link referenced by those scripts has become
 13inaccessible (Google reports the file as "not accessible", i.e. its sharing permission was revoked or
 14changed, not a rate limit). An older, still-accessible mirror of the same source exists at a different
 15Google Drive file id, but the file served from that mirror is truncated: it only contains the "AMD
 16Part1", "AMD Part2", "Normal Part1" and "Normal Part2" categories in full (1269 of 1672 manually graded
 17images, all complete for these four categories), with the "DME" category entirely missing (0 of 403).
 18The truncated archive's own zip central directory is also corrupted (its end-of-central-directory record
 19does not reflect the true file length), so this module recovers the contained files by scanning for raw
 20zip local file headers directly rather than relying on the central directory. This is real, verified
 21recovered data (checked against the dataset's own path manifest, `manual_paths.csv`), not a workaround
 22of any access restriction: the mirror itself has no password beyond the one already published for the
 23primary source, and no additional data exists beyond what is recovered here. Consequently, this loader
 24only supports the "AMD Part1", "AMD Part2", "Normal Part1" and "Normal Part2" categories; "DME" is not
 25available. This should be treated as a known, permanent limitation until the primary source is restored.
 26
 27NOTE: The masks are index-encoded label maps (not the RGB visualization masks the archive also ships)
 28with values 0-5: 0 is background, and 1-5 correspond to the 5 retinal layer bands delimited by the ILM,
 29OPL-Henles, IS/OS junction, IBRPE and OBRPE boundaries (per the dataset's own README), in order from the
 30vitreous side down. The raw OCT images and their masks are not the same resolution in the original
 31archive: masks are rendered on a fixed 512x512 canvas, while raw images keep their native (smaller) width.
 32This loader resizes the mask to the raw image's shape with nearest-neighbor interpolation (verified to
 33align the layer bands with the visible retinal structure in the raw image, not merely equal in aspect).
 34"""
 35
 36import os
 37import csv
 38import struct
 39import zlib
 40from glob import glob
 41from natsort import natsorted
 42from typing import Union, Literal, Tuple, List
 43
 44import numpy as np
 45import imageio.v3 as imageio
 46from PIL import Image
 47
 48from torch.utils.data import Dataset, DataLoader
 49
 50import torch_em
 51
 52from .. import util
 53
 54
 55MASKS_URL = "https://ndownloader.figshare.com/files/44436359"
 56MASKS_CHECKSUM = "ae61b43b11c51f39a755b3b63ba4a6353adb99b1e581792c360caa6740df2611"
 57
 58RAW_IMAGES_URL = "https://drive.google.com/uc?id=1y7yKlDR4sP8bJ-_updZFBtHD6Yq_FO03"
 59RAW_IMAGES_CHECKSUM = "da9229443c0601c6ed66d67ceebe39c5edc78762ec5e54d4917105ee635a5d11"
 60RAW_IMAGES_PASSWORD = b"MCME2017"
 61
 62CATEGORIES = ["AMD Part1", "AMD Part2", "Normal Part1", "Normal Part2"]
 63
 64
 65def _decrypt_zipcrypto(data, password):
 66    key0, key1, key2 = 305419896, 591751049, 878082192
 67
 68    def crc32_step(crc, byte):
 69        return zlib.crc32(bytes([byte]), crc ^ 0xffffffff) ^ 0xffffffff
 70
 71    def update_keys(byte):
 72        nonlocal key0, key1, key2
 73        key0 = crc32_step(key0, byte)
 74        key1 = (key1 + (key0 & 0xff)) & 0xffffffff
 75        key1 = (key1 * 134775813 + 1) & 0xffffffff
 76        key2 = crc32_step(key2, (key1 >> 24) & 0xff)
 77
 78    for byte in password:
 79        update_keys(byte)
 80
 81    out = bytearray()
 82    for byte in data:
 83        temp = (key2 | 2) & 0xffff
 84        decrypt_byte = ((temp * (temp ^ 1)) >> 8) & 0xff
 85        plain = byte ^ decrypt_byte
 86        out.append(plain)
 87        update_keys(plain)
 88    return bytes(out)
 89
 90
 91def _carve_zip_entries(zip_path, dst):
 92    # The archive at `RAW_IMAGES_URL` is truncated and its central directory does not reflect the
 93    # true file content, so standard zip readers (including Python's own 'zipfile') refuse to open
 94    # it. This recovers every intact entry by scanning for local file header signatures directly.
 95    local_header_sig = b"PK\x03\x04"
 96    data_descriptor_sig = b"PK\x07\x08"
 97
 98    with open(zip_path, "rb") as f:
 99        data = f.read()
100
101    pos = 0
102    while True:
103        idx = data.find(local_header_sig, pos)
104        if idx == -1:
105            break
106
107        header = data[idx:idx + 30]
108        if len(header) < 30:
109            break
110
111        _, _, flags, method, _, _, _, comp_size, uncomp_size, name_len, extra_len = struct.unpack(
112            "<IHHHHHIIIHH", header
113        )
114        name_start = idx + 30
115        name = data[name_start:name_start + name_len].decode("utf-8", errors="replace")
116        data_start = name_start + name_len + extra_len
117
118        has_data_descriptor = bool(flags & 0x08)
119        is_encrypted = bool(flags & 0x01)
120
121        if has_data_descriptor and comp_size == 0:
122            dd_idx = data.find(data_descriptor_sig, data_start)
123            if dd_idx == -1:
124                pos = idx + 4
125                continue
126            comp_data = data[data_start:dd_idx]
127            _, _, comp_size, uncomp_size = struct.unpack("<IIII", data[dd_idx:dd_idx + 16])
128            next_pos = dd_idx + 16
129        else:
130            comp_data = data[data_start:data_start + comp_size]
131            next_pos = data_start + comp_size
132
133        if name and not name.endswith("/") and uncomp_size > 0:
134            raw = comp_data
135            if is_encrypted:
136                raw = _decrypt_zipcrypto(comp_data, RAW_IMAGES_PASSWORD)[12:]
137
138            try:
139                if method == 8:
140                    content = zlib.decompressobj(-15).decompress(raw)
141                elif method == 0:
142                    content = raw[:uncomp_size]
143                else:
144                    content = None
145            except zlib.error:
146                content = None
147
148            if content is not None:
149                out_path = os.path.join(dst, name)
150                os.makedirs(os.path.dirname(out_path), exist_ok=True)
151                with open(out_path, "wb") as out_f:
152                    out_f.write(content)
153
154        pos = next_pos
155
156
157def get_oct5k_data(path: Union[os.PathLike, str], download: bool = False) -> str:
158    """Download the OCT5k dataset.
159
160    Args:
161        path: Filepath to a folder where the data is downloaded for further processing.
162        download: Whether to download the data if it is not present.
163
164    Returns:
165        Filepath where the data is downloaded.
166    """
167    os.makedirs(path, exist_ok=True)
168
169    masks_dir = os.path.join(path, "OCT5k")
170    if not os.path.exists(masks_dir):
171        masks_zip = os.path.join(path, "OCT5k.zip")
172        util.download_source(path=masks_zip, url=MASKS_URL, download=download, checksum=MASKS_CHECKSUM)
173        util.unzip(zip_path=masks_zip, dst=path, remove=False)
174
175    raw_dir = os.path.join(path, "raw")
176    if not glob(os.path.join(raw_dir, "**", "*.TIFF"), recursive=True):
177        raw_zip = os.path.join(path, "rasti_old.zip")
178        util.download_source_gdrive(path=raw_zip, url=RAW_IMAGES_URL, download=download, checksum=RAW_IMAGES_CHECKSUM)
179        _carve_zip_entries(raw_zip, raw_dir)
180
181    return path
182
183
184def get_oct5k_paths(
185    path: Union[os.PathLike, str],
186    grading: Literal["1", "2", "3"] = "1",
187    download: bool = False,
188) -> Tuple[List[str], List[str]]:
189    """Get paths to the OCT5k data.
190
191    Args:
192        path: Filepath to a folder where the data is downloaded for further processing.
193        grading: The choice of manual grader ('1', '2' or '3'). All three graded every image.
194        download: Whether to download the data if it is not present.
195
196    Returns:
197        List of filepaths for the image data.
198        List of filepaths for the label data.
199    """
200    data_dir = get_oct5k_data(path, download)
201
202    manual_paths_csv = os.path.join(data_dir, "OCT5k", "Scripts", "paths", "manual_paths.csv")
203    masks_root = os.path.join(data_dir, "OCT5k", "Masks", "Masks_Manual", f"Grading_{grading}")
204    raw_root = os.path.join(data_dir, "raw")
205
206    resized_masks_dir = os.path.join(data_dir, f"masks_resized_grading_{grading}")
207
208    image_paths, gt_paths = [], []
209    with open(manual_paths_csv) as f:
210        for row in csv.reader(f):
211            mask_rel_path = row[0].replace("../Images/Images_Manual/", "")
212            raw_rel_path = row[1].replace("./Macular-Dataset-R.Rasti_old/", "")
213
214            raw_path = os.path.join(raw_root, raw_rel_path)
215            if not os.path.exists(raw_path):
216                # This image belongs to the 'DME' category, or otherwise was not among the files
217                # recoverable from the truncated archive. See the module docstring.
218                continue
219
220            mask_path = os.path.join(masks_root, mask_rel_path)
221            resized_mask_path = os.path.join(resized_masks_dir, mask_rel_path)
222            if not os.path.exists(resized_mask_path):
223                raw_size = imageio.imread(raw_path).shape[::-1]  # (H, W) -> (W, H) for PIL 'resize'.
224                mask = Image.open(mask_path)
225                mask_resized = mask.resize(raw_size, Image.NEAREST)
226                os.makedirs(os.path.dirname(resized_mask_path), exist_ok=True)
227                imageio.imwrite(resized_mask_path, np.array(mask_resized))
228
229            image_paths.append(raw_path)
230            gt_paths.append(resized_mask_path)
231
232    image_paths, gt_paths = natsorted(image_paths), natsorted(gt_paths)
233
234    assert len(image_paths) == len(gt_paths) and len(image_paths) > 0
235
236    return image_paths, gt_paths
237
238
239def get_oct5k_dataset(
240    path: Union[os.PathLike, str],
241    patch_shape: Tuple[int, int],
242    grading: Literal["1", "2", "3"] = "1",
243    resize_inputs: bool = False,
244    download: bool = False,
245    **kwargs
246) -> Dataset:
247    """Get the OCT5k dataset for retinal layer band segmentation in OCT B-scans.
248
249    Args:
250        path: Filepath to a folder where the data is downloaded for further processing.
251        patch_shape: The patch shape to use for training.
252        grading: The choice of manual grader ('1', '2' or '3'). All three graded every image.
253        resize_inputs: Whether to resize inputs to the desired patch shape.
254        download: Whether to download the data if it is not present.
255        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
256
257    Returns:
258        The segmentation dataset.
259    """
260    image_paths, gt_paths = get_oct5k_paths(path, grading, download)
261
262    if resize_inputs:
263        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
264        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
265            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
266        )
267
268    return torch_em.default_segmentation_dataset(
269        raw_paths=image_paths,
270        raw_key=None,
271        label_paths=gt_paths,
272        label_key=None,
273        patch_shape=patch_shape,
274        is_seg_dataset=False,
275        **kwargs
276    )
277
278
279def get_oct5k_loader(
280    path: Union[os.PathLike, str],
281    batch_size: int,
282    patch_shape: Tuple[int, int],
283    grading: Literal["1", "2", "3"] = "1",
284    resize_inputs: bool = False,
285    download: bool = False,
286    **kwargs
287) -> DataLoader:
288    """Get the OCT5k dataloader for retinal layer band segmentation in OCT B-scans.
289
290    Args:
291        path: Filepath to a folder where the data is downloaded for further processing.
292        batch_size: The batch size for training.
293        patch_shape: The patch shape to use for training.
294        grading: The choice of manual grader ('1', '2' or '3'). All three graded every image.
295        resize_inputs: Whether to resize inputs to the desired patch shape.
296        download: Whether to download the data if it is not present.
297        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
298
299    Returns:
300        The DataLoader.
301    """
302    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
303    dataset = get_oct5k_dataset(path, patch_shape, grading, resize_inputs, download, **ds_kwargs)
304    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)