torch_em.data.datasets.light_microscopy.isbi14

  1"""The ISBI14 dataset contains cervical cytology images with nucleus and cytoplasm
  2instance segmentation annotations, released for the Overlapping Cervical Cytology Image
  3Segmentation Challenge at ISBI 2014.
  4
  5The dataset has two image sources:
  6- 'synthetic': 945 synthetic Pap-smear images (45 train, 900 test) composed from the real
  7  images below, with per-cell nucleus and cytoplasm instance masks.
  8- 'real': 16 real extended-depth-of-field (EDF) Pap-smear images with nucleus instance masks.
  9  No cytoplasm annotation is available for the real images.
 10
 11NOTE: 4 of the 16 real images (EDF000-EDF003) ship ground truth as hundreds of few-pixel
 12fragments instead of per-cell nucleus masks. This loader drops them and only exposes the
 13remaining 12 real images.
 14
 15NOTE: No data license is published on the challenge site, its linked pages, or inside the
 16released archives.
 17
 18The dataset is located at https://cs.adelaide.edu.au/~carneiro/isbi14_challenge/dataset.html.
 19This dataset is from the publications https://doi.org/10.1109/JBHI.2016.2601609 and
 20https://doi.org/10.1109/TIP.2015.2389619. Please cite them if you use this dataset in your research.
 21"""
 22
 23import os
 24import ssl
 25from glob import glob
 26from natsort import natsorted
 27from typing import List, Literal, Tuple, Union
 28
 29import numpy as np
 30import requests
 31import h5py
 32import scipy.io as sio
 33import imageio.v3 as imageio
 34from scipy import ndimage
 35from requests.adapters import HTTPAdapter
 36from tqdm import tqdm
 37
 38from torch.utils.data import Dataset, DataLoader
 39
 40import torch_em
 41
 42from .. import util
 43
 44
 45URL = "https://cs.adelaide.edu.au/~carneiro/isbi14_challenge/Dataset.zip"
 46CHECKSUM = "618e386cd87722ad0bca5ebe8570fc7d96ba4cf011eb8b018f865c4e7d8328f0"
 47
 48# The challenge server only supports legacy TLS renegotiation, which OpenSSL 3 disables by default.
 49SSL_OP_LEGACY_SERVER_CONNECT = 0x4
 50
 51# These real images ship ground truth as hundreds of few-pixel fragments, not usable nucleus masks.
 52BROKEN_REAL_IMAGES = ("EDF000", "EDF001", "EDF002", "EDF003")
 53
 54IMAGE_SOURCES = ("synthetic", "real")
 55LABEL_CHOICES = ("nucleus", "cytoplasm")
 56
 57
 58class _LegacyRenegotiationAdapter(HTTPAdapter):
 59    """Allow the TLS renegotiation that the challenge server still requires."""
 60
 61    def init_poolmanager(self, *args, **kwargs):
 62        context = ssl.create_default_context()
 63        context.options |= SSL_OP_LEGACY_SERVER_CONNECT
 64        kwargs["ssl_context"] = context
 65        return super().init_poolmanager(*args, **kwargs)
 66
 67
 68def _download_zip(zip_path: str, download: bool) -> None:
 69    if os.path.exists(zip_path):
 70        return
 71    if not download:
 72        raise RuntimeError(f"Cannot find the data at {zip_path}, but download was set to False")
 73
 74    session = requests.Session()
 75    session.mount("https://cs.adelaide.edu.au", _LegacyRenegotiationAdapter())
 76
 77    tmp_path = f"{zip_path}.incomplete"
 78    with session.get(URL, stream=True, timeout=60) as r:
 79        r.raise_for_status()
 80        file_size = int(r.headers.get("Content-Length", 0))
 81        with tqdm.wrapattr(r.raw, "read", total=file_size, desc=f"Download {URL}") as r_raw, open(tmp_path, "wb") as f:
 82            for chunk in iter(lambda: r_raw.read(1024 * 1024), b""):
 83                f.write(chunk)
 84
 85    this_checksum = util.get_checksum(tmp_path)
 86    if this_checksum != CHECKSUM:
 87        raise RuntimeError(f"The checksum of the download does not match. Expected: {CHECKSUM}, got: {this_checksum}")
 88    os.replace(tmp_path, zip_path)
 89
 90
 91def get_isbi14_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 92    """Download the ISBI14 dataset.
 93
 94    Args:
 95        path: Filepath to a folder where the downloaded data will be saved.
 96        download: Whether to download the data if it is not present.
 97
 98    Returns:
 99        The filepath to the extracted data directory.
100    """
101    data_dir = os.path.join(path, "Dataset")
102    if os.path.exists(data_dir):
103        return data_dir
104
105    os.makedirs(path, exist_ok=True)
106    zip_path = os.path.join(path, "Dataset.zip")
107    _download_zip(zip_path, download)
108    util.unzip(zip_path, path)
109
110    return data_dir
111
112
113def _load_mat_array(f_or_dict, key, index, is_h5):
114    """Load one entry of a MATLAB cell array, from either a scipy or an h5py backed file."""
115    if is_h5:
116        return np.array(f_or_dict[f_or_dict[key][0, index]][()]).T
117    return f_or_dict[key][index, 0]
118
119
120def _load_cytoplasm_entry(f_or_dict, key, index, is_h5):
121    """Load the per-cell cytoplasm masks of one image, from either a scipy or an h5py backed file."""
122    if is_h5:
123        refs = np.array(f_or_dict[f_or_dict[key][0, index]][()]).ravel()
124        return [np.array(f_or_dict[ref][()]).T for ref in refs]
125    entry = f_or_dict[key][index, 0]
126    return [entry[j, 0] for j in range(entry.shape[0])]
127
128
129def _cytoplasm_to_labels(cytoplasm_masks: List[np.ndarray]) -> np.ndarray:
130    """Merge per-cell cytoplasm masks into one label image, larger cells first so a smaller,
131    more visible cell on top keeps its own label where cytoplasms overlap."""
132    ordered = sorted(range(len(cytoplasm_masks)), key=lambda i: int(cytoplasm_masks[i].sum()), reverse=True)
133    labels = np.zeros_like(cytoplasm_masks[0], dtype="uint16")
134    for instance_id, i in enumerate(ordered, start=1):
135        labels[cytoplasm_masks[i] > 0] = instance_id
136    return labels
137
138
139def _prepare_synthetic_split(data_dir: str, split: Literal["train", "test"]) -> Tuple[str, str, str]:
140    raw_dir = os.path.join(data_dir, "synthetic_preprocessed", split, "raw")
141    nucleus_dir = os.path.join(data_dir, "synthetic_preprocessed", split, "labels_nucleus")
142    cytoplasm_dir = os.path.join(data_dir, "synthetic_preprocessed", split, "labels_cytoplasm")
143
144    raw_mat = os.path.join(data_dir, "Synthetic", f"{split}set.mat")
145    gt_mat = os.path.join(data_dir, "Synthetic", f"{split}set_GT.mat")
146
147    try:
148        gt = sio.loadmat(gt_mat)
149        gt_is_h5 = False
150    except NotImplementedError:
151        gt = h5py.File(gt_mat, "r")
152        gt_is_h5 = True
153
154    n_images = gt[f"{split}_Nuclei"].shape[0 if not gt_is_h5 else 1]
155
156    is_cached = all(
157        os.path.exists(d) and len(glob(os.path.join(d, "*.tif"))) == n_images
158        for d in (raw_dir, nucleus_dir, cytoplasm_dir)
159    )
160    if is_cached:
161        if gt_is_h5:
162            gt.close()
163        return raw_dir, nucleus_dir, cytoplasm_dir
164
165    for d in (raw_dir, nucleus_dir, cytoplasm_dir):
166        os.makedirs(d, exist_ok=True)
167
168    with h5py.File(raw_mat, "r") as raw_f:
169        for i in tqdm(range(n_images), desc=f"Preprocess ISBI14 synthetic {split} images"):
170            name = f"{i:04d}.tif"
171
172            raw = _load_mat_array(raw_f, f"{split}set", i, is_h5=True)
173            imageio.imwrite(os.path.join(raw_dir, name), raw)
174
175            nucleus_mask = _load_mat_array(gt, f"{split}_Nuclei", i, gt_is_h5)
176            nucleus_labels, _ = ndimage.label(nucleus_mask)
177            imageio.imwrite(os.path.join(nucleus_dir, name), nucleus_labels.astype("uint16"))
178
179            cytoplasm_masks = _load_cytoplasm_entry(gt, f"{split}_Cytoplasm", i, gt_is_h5)
180            cytoplasm_labels = _cytoplasm_to_labels(cytoplasm_masks)
181            imageio.imwrite(os.path.join(cytoplasm_dir, name), cytoplasm_labels)
182
183    if gt_is_h5:
184        gt.close()
185
186    return raw_dir, nucleus_dir, cytoplasm_dir
187
188
189def _prepare_real(data_dir: str) -> Tuple[str, str]:
190    raw_dir = os.path.join(data_dir, "real_preprocessed", "raw")
191    label_dir = os.path.join(data_dir, "real_preprocessed", "labels_nucleus")
192
193    image_paths = [
194        p for p in natsorted(glob(os.path.join(data_dir, "EDF", "*.png")))
195        if not os.path.basename(p).endswith("_GT.png")
196        and os.path.splitext(os.path.basename(p))[0] not in BROKEN_REAL_IMAGES
197    ]
198    is_cached = all(
199        os.path.exists(d) and len(glob(os.path.join(d, "*.tif"))) == len(image_paths)
200        for d in (raw_dir, label_dir)
201    )
202    if is_cached:
203        return raw_dir, label_dir
204
205    os.makedirs(raw_dir, exist_ok=True)
206    os.makedirs(label_dir, exist_ok=True)
207
208    for image_path in tqdm(image_paths, desc="Preprocess ISBI14 real images"):
209        name = os.path.splitext(os.path.basename(image_path))[0]
210        gt_path = os.path.join(data_dir, "EDF", f"{name}_GT.png")
211
212        raw = imageio.imread(image_path)
213        mask = imageio.imread(gt_path) > 0
214        labels, _ = ndimage.label(mask)
215
216        imageio.imwrite(os.path.join(raw_dir, f"{name}.tif"), raw)
217        imageio.imwrite(os.path.join(label_dir, f"{name}.tif"), labels.astype("uint16"))
218
219    return raw_dir, label_dir
220
221
222def get_isbi14_paths(
223    path: Union[os.PathLike, str],
224    image_source: Literal["synthetic", "real"] = "synthetic",
225    label_choice: Literal["nucleus", "cytoplasm"] = "nucleus",
226    split: Literal["train", "test"] = "train",
227    download: bool = False,
228) -> Tuple[List[str], List[str]]:
229    """Get paths to the ISBI14 data.
230
231    Args:
232        path: Filepath to a folder where the downloaded data will be saved.
233        image_source: The image source. Either 'synthetic' or 'real'.
234        label_choice: The segmentation target. Either 'nucleus' or 'cytoplasm'.
235            No cytoplasm annotation exists for the real images.
236        split: The data split. Only used for the synthetic images.
237        download: Whether to download the data if it is not present.
238
239    Returns:
240        List of filepaths for the image data.
241        List of filepaths for the label data.
242    """
243    if image_source not in IMAGE_SOURCES:
244        raise ValueError(f"'{image_source}' is not a valid image source. Choose from {list(IMAGE_SOURCES)}.")
245    if label_choice not in LABEL_CHOICES:
246        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose from {list(LABEL_CHOICES)}.")
247    if image_source == "real" and label_choice == "cytoplasm":
248        raise ValueError("No cytoplasm annotation exists for the real images. Use label_choice='nucleus'.")
249
250    data_dir = get_isbi14_data(path, download)
251
252    if image_source == "real":
253        raw_dir, label_dir = _prepare_real(data_dir)
254    else:
255        raw_dir, nucleus_dir, cytoplasm_dir = _prepare_synthetic_split(data_dir, split)
256        label_dir = nucleus_dir if label_choice == "nucleus" else cytoplasm_dir
257
258    raw_paths = natsorted(glob(os.path.join(raw_dir, "*.tif")))
259    label_paths = natsorted(glob(os.path.join(label_dir, "*.tif")))
260
261    return raw_paths, label_paths
262
263
264def get_isbi14_dataset(
265    path: Union[os.PathLike, str],
266    patch_shape: Tuple[int, ...],
267    image_source: Literal["synthetic", "real"] = "synthetic",
268    label_choice: Literal["nucleus", "cytoplasm"] = "nucleus",
269    split: Literal["train", "test"] = "train",
270    download: bool = False,
271    **kwargs,
272) -> Dataset:
273    """Get the ISBI14 dataset for cervical cell instance segmentation.
274
275    Args:
276        path: Filepath to a folder where the downloaded data will be saved.
277        patch_shape: The patch shape to use for training.
278        image_source: The image source. Either 'synthetic' or 'real'.
279        label_choice: The segmentation target. Either 'nucleus' or 'cytoplasm'.
280            No cytoplasm annotation exists for the real images.
281        split: The data split. Only used for the synthetic images.
282        download: Whether to download the data if it is not present.
283        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
284
285    Returns:
286        The segmentation dataset.
287    """
288    raw_paths, label_paths = get_isbi14_paths(path, image_source, label_choice, split, download)
289
290    return torch_em.default_segmentation_dataset(
291        raw_paths=raw_paths,
292        raw_key=None,
293        label_paths=label_paths,
294        label_key=None,
295        patch_shape=patch_shape,
296        **kwargs,
297    )
298
299
300def get_isbi14_loader(
301    path: Union[os.PathLike, str],
302    batch_size: int,
303    patch_shape: Tuple[int, ...],
304    image_source: Literal["synthetic", "real"] = "synthetic",
305    label_choice: Literal["nucleus", "cytoplasm"] = "nucleus",
306    split: Literal["train", "test"] = "train",
307    download: bool = False,
308    **kwargs,
309) -> DataLoader:
310    """Get the ISBI14 dataloader for cervical cell instance segmentation.
311
312    Args:
313        path: Filepath to a folder where the downloaded data will be saved.
314        batch_size: The batch size for training.
315        patch_shape: The patch shape to use for training.
316        image_source: The image source. Either 'synthetic' or 'real'.
317        label_choice: The segmentation target. Either 'nucleus' or 'cytoplasm'.
318            No cytoplasm annotation exists for the real images.
319        split: The data split. Only used for the synthetic images.
320        download: Whether to download the data if it is not present.
321        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
322
323    Returns:
324        The DataLoader.
325    """
326    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
327    dataset = get_isbi14_dataset(path, patch_shape, image_source, label_choice, split, download, **ds_kwargs)
328    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)