torch_em.data.datasets.medical.retouch

  1"""The RETOUCH dataset contains annotations for intraretinal fluid (IRF), subretinal fluid (SRF)
  2and pigment epithelial detachment (PED) segmentation in 3D retinal OCT volumes.
  3
  4The dataset was curated for the RETOUCH challenge (MICCAI 2017, https://retouch.grand-challenge.org).
  5It comprises 112 OCT volumes acquired with three different device manufacturers (Cirrus, Spectralis
  6and Topcon), out of which 70 volumes (24 Cirrus, 24 Spectralis, 22 Topcon) form the labeled training
  7set and 42 volumes form the unlabeled test set. Only the labeled training set is exposed here, as the
  8test set annotations are not publicly available.
  9
 10NOTE: The dataset is gated and cannot be downloaded automatically by 'torch_em'. To obtain it:
 11- Visit https://retouch.grand-challenge.org and click on 'Join' to register for the challenge.
 12- Complete the registration form and sign the 'RETOUCH-Agreement_of_Data_Confidentiality' document.
 13- Once your registration is approved (contact hrvoje.bogunovic@meduniwien.ac.at for questions),
 14  you will get access to the 'Download' page, from which the training data must be downloaded.
 15- Extract the downloaded archive(s) to the desired 'path'. The expected layout per vendor is:
 16  '<path>/RETOUCH-TrainingSet-<vendor>/TRAIN<id>/oct.mhd', 'oct.raw', 'reference.mhd', 'reference.raw',
 17  where '<vendor>' is one of 'Cirrus', 'Spectralis', 'Topcon'.
 18
 19The original MetaImage ('.mhd' / '.raw') volumes are converted to hdf5 once and then loaded from there.
 20
 21The reference annotations use the label ids: 0 = background, 1 = IRF, 2 = SRF, 3 = PED.
 22
 23The dataset is from the publication https://doi.org/10.1109/TMI.2019.2901398.
 24Please cite it if you use this dataset in your research.
 25"""
 26
 27import os
 28from glob import glob
 29from tqdm import tqdm
 30from natsort import natsorted
 31from typing import Union, Tuple, List, Literal, Sequence
 32
 33import numpy as np
 34
 35from torch.utils.data import Dataset, DataLoader
 36
 37import torch_em
 38
 39from .. import util
 40
 41
 42VENDORS = ["Cirrus", "Spectralis", "Topcon"]
 43
 44METAIMAGE_DTYPES = {
 45    "MET_CHAR": "int8", "MET_UCHAR": "uint8", "MET_SHORT": "int16", "MET_USHORT": "uint16",
 46    "MET_INT": "int32", "MET_UINT": "uint32", "MET_FLOAT": "float32", "MET_DOUBLE": "float64",
 47}
 48
 49
 50def _read_metaimage(path):
 51    """Read an uncompressed MetaImage ('.mhd' + '.raw') volume with numpy, returning it in 'zyx' order."""
 52    header = {}
 53    with open(path, "r") as f:
 54        for line in f:
 55            if "=" in line:
 56                key, value = line.split("=", 1)
 57                header[key.strip()] = value.strip()
 58
 59    if header.get("CompressedData", "False") == "True":
 60        raise NotImplementedError(f"Compressed MetaImage data is not supported: '{path}'.")
 61
 62    shape = tuple(int(s) for s in header["DimSize"].split())[::-1]
 63    dtype = np.dtype(METAIMAGE_DTYPES[header["ElementType"]])
 64    if header.get("BinaryDataByteOrderMSB", "False") == "True":
 65        dtype = dtype.newbyteorder(">")
 66
 67    raw_path = os.path.join(os.path.dirname(path), header["ElementDataFile"])
 68    return np.fromfile(raw_path, dtype=dtype).reshape(shape)
 69
 70
 71def get_retouch_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 72    """Obtain the RETOUCH dataset.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        Filepath to the folder with the converted hdf5 volumes.
 80    """
 81    data_dir = os.path.join(path, "preprocessed")
 82    if os.path.exists(data_dir):
 83        return data_dir
 84
 85    vendor_dirs = [os.path.join(path, f"RETOUCH-TrainingSet-{vendor}") for vendor in VENDORS]
 86    if not any(os.path.exists(vendor_dir) for vendor_dir in vendor_dirs):
 87        msg = "The RETOUCH dataset is gated and cannot be downloaded automatically. To obtain it:\n" \
 88            "- Visit https://retouch.grand-challenge.org and click on 'Join' to register for the challenge.\n" \
 89            "- Complete the registration form and sign the 'RETOUCH-Agreement_of_Data_Confidentiality' " \
 90            "document. Contact hrvoje.bogunovic@meduniwien.ac.at for registration questions.\n" \
 91            "- Once approved, download the training data from the challenge 'Download' page.\n" \
 92            f"- Extract the downloaded archive(s) so that '{path}' contains one folder per vendor, " \
 93            "named 'RETOUCH-TrainingSet-Cirrus', 'RETOUCH-TrainingSet-Spectralis' and " \
 94            "'RETOUCH-TrainingSet-Topcon', each with 'TRAIN<id>' subfolders containing " \
 95            "'oct.mhd', 'oct.raw', 'reference.mhd' and 'reference.raw'."
 96        if download:
 97            msg = "Download is set to True, but 'torch_em' cannot download this dataset.\n" + msg
 98        raise RuntimeError(msg)
 99
100    import h5py
101
102    os.makedirs(data_dir, exist_ok=True)
103    for vendor_dir in vendor_dirs:
104        if not os.path.exists(vendor_dir):
105            continue
106
107        vendor = os.path.basename(vendor_dir).replace("RETOUCH-TrainingSet-", "")
108        case_dirs = natsorted(glob(os.path.join(vendor_dir, "TRAIN*")))
109        for case_dir in tqdm(case_dirs, desc=f"Converting RETOUCH volumes for '{os.path.basename(vendor_dir)}'"):
110            case_id = os.path.basename(case_dir)
111            out_path = os.path.join(data_dir, f"{vendor}_{case_id}.h5")
112            if os.path.exists(out_path):
113                continue
114
115            raw = _read_metaimage(os.path.join(case_dir, "oct.mhd"))
116            labels = _read_metaimage(os.path.join(case_dir, "reference.mhd")).astype("uint8")
117            assert raw.shape == labels.shape, f"Shape mismatch for '{case_dir}': {raw.shape} vs. {labels.shape}."
118
119            with h5py.File(out_path, "w") as f:
120                f.create_dataset("raw", data=raw, compression="gzip")
121                f.create_dataset("labels", data=labels, compression="gzip")
122
123    return data_dir
124
125
126def get_retouch_paths(
127    path: Union[os.PathLike, str], vendor: Union[Literal["Cirrus", "Spectralis", "Topcon"], Sequence[str], None] = None,
128    download: bool = False,
129) -> List[str]:
130    """Get paths to the RETOUCH data.
131
132    Args:
133        path: Filepath to a folder where the data is downloaded for further processing.
134        vendor: The choice of device vendor(s). By default, volumes from all vendors are used.
135        download: Whether to download the data if it is not present.
136
137    Returns:
138        List of filepaths for the hdf5 volumes, which store the image data at 'raw' and the label data at 'labels'.
139    """
140    data_dir = get_retouch_data(path, download)
141
142    if vendor is None:
143        vendors = VENDORS
144    else:
145        vendors = [vendor] if isinstance(vendor, str) else list(vendor)
146        for v in vendors:
147            if v not in VENDORS:
148                raise ValueError(f"'{v}' is not a valid vendor. Choose from {VENDORS}.")
149
150    volume_paths = []
151    for v in vendors:
152        volume_paths.extend(natsorted(glob(os.path.join(data_dir, f"{v}_TRAIN*.h5"))))
153
154    assert len(volume_paths) > 0, f"No RETOUCH volumes found at '{data_dir}'."
155    return volume_paths
156
157
158def get_retouch_dataset(
159    path: Union[os.PathLike, str],
160    patch_shape: Tuple[int, ...],
161    vendor: Union[Literal["Cirrus", "Spectralis", "Topcon"], Sequence[str], None] = None,
162    resize_inputs: bool = False,
163    download: bool = False,
164    **kwargs
165) -> Dataset:
166    """Get the RETOUCH dataset for retinal fluid segmentation.
167
168    Args:
169        path: Filepath to a folder where the data is downloaded for further processing.
170        patch_shape: The patch shape to use for training.
171        vendor: The choice of device vendor(s). By default, volumes from all vendors are used.
172        resize_inputs: Whether to resize inputs to the desired patch shape.
173        download: Whether to download the data if it is not present.
174        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
175
176    Returns:
177        The segmentation dataset.
178    """
179    volume_paths = get_retouch_paths(path, vendor, download)
180
181    if resize_inputs:
182        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
183        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
184            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
185        )
186
187    return torch_em.default_segmentation_dataset(
188        raw_paths=volume_paths,
189        raw_key="raw",
190        label_paths=volume_paths,
191        label_key="labels",
192        patch_shape=patch_shape,
193        is_seg_dataset=True,
194        **kwargs
195    )
196
197
198def get_retouch_loader(
199    path: Union[os.PathLike, str],
200    batch_size: int,
201    patch_shape: Tuple[int, ...],
202    vendor: Union[Literal["Cirrus", "Spectralis", "Topcon"], Sequence[str], None] = None,
203    resize_inputs: bool = False,
204    download: bool = False,
205    **kwargs
206) -> DataLoader:
207    """Get the RETOUCH dataloader for retinal fluid segmentation.
208
209    Args:
210        path: Filepath to a folder where the data is downloaded for further processing.
211        batch_size: The batch size for training.
212        patch_shape: The patch shape to use for training.
213        vendor: The choice of device vendor(s). By default, volumes from all vendors are used.
214        resize_inputs: Whether to resize inputs to the desired patch shape.
215        download: Whether to download the data if it is not present.
216        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
217
218    Returns:
219        The DataLoader.
220    """
221    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
222    dataset = get_retouch_dataset(path, patch_shape, vendor, resize_inputs, download, **ds_kwargs)
223    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)