torch_em.data.datasets.light_microscopy.micro_bench

  1"""Micro-Bench contains microscopy images for vision-language understanding.
  2
  3This loader exposes polygon annotations from Burgess et al., CellCognition, OpenCell and Wu et al.
  4The dataset is located at https://huggingface.co/datasets/jnirschl/uBench and is from the publication
  5https://doi.org/10.52202/079017-0965. Please cite it if you use this dataset in your research.
  6"""
  7
  8import os
  9from glob import glob
 10from io import BytesIO
 11from pathlib import Path
 12from typing import List, Literal, Optional, Tuple, Union
 13
 14import numpy as np
 15
 16from torch.utils.data import DataLoader, Dataset
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23REVISION = "3f2c5b590bc7a208d5b60f3527ce4c76a331aa2b"
 24URL = (
 25    "https://huggingface.co/datasets/jnirschl/uBench/resolve/"
 26    f"{REVISION}/perception/0.1.0/{{filename}}?download=true"
 27)
 28
 29FILES = {
 30    "ubench-test-00000-of-00007.arrow": "3d346134e6dffed41f74da64c1c173f5efd3a10200805796a030357297b630f1",
 31    "ubench-test-00001-of-00007.arrow": "790591636ae679afa1161664d207c2bb2e3d583690d03967e38364ba0f31ccac",
 32    "ubench-test-00002-of-00007.arrow": "c370cd84e6ed9b00483ef3d20f5c06d5ff2264733ea87f2df98cf9a22820bf74",
 33    "ubench-test-00003-of-00007.arrow": "832cdd6f73a132028584b4e47ec1afe643e0490dd8ec9686eabbb9de3eb7f814",
 34    "ubench-test-00004-of-00007.arrow": "e2741415cef1327ba243f21000ac43aab2d36b61a9c2c862a1d0e21f200f4a69",
 35    "ubench-test-00005-of-00007.arrow": "89deae84ed2dab1bb1fc40be0aefe81b8d94590cec548d1950ff5ff853143d34",
 36    "ubench-test-00006-of-00007.arrow": "534beca2084a250d4436340c78accc9fe7cc2d6063d4884302e7e7f141edad6c",
 37}
 38
 39SOURCES = {
 40    "synthetic_cells": ("burgess", "cell"),
 41    "synthetic_nuclei": ("burgess", "nucleus"),
 42    "fluo_nuclei": ("cellcognition", "instances"),
 43    "protein_localization_nuclei": ("opencell", "instances"),
 44    "em_mitochondria": ("wu", "instances"),
 45}
 46
 47
 48def _get_source(row) -> Optional[str]:
 49    dataset = row["dataset"]
 50    if dataset is not None and dataset.startswith("burgess_et_al_2024_"):
 51        return "burgess"
 52    if dataset == "opencell":
 53        return "opencell"
 54    if dataset == "wu_et_al_2023":
 55        return "wu"
 56
 57    classes = {annotation["className"] for annotation in row["polygon"] or []}
 58    if dataset is None and "H2B-mCherry" in classes:
 59        return "cellcognition"
 60    return None
 61
 62
 63def _get_sample_id(row, source: str) -> str:
 64    dataset = source if row["dataset"] is None else row["dataset"]
 65    return f"{dataset}_{row['image_id']}"
 66
 67
 68def _rasterize(polygons, shape: Tuple[int, int], label_choice: Optional[str] = None) -> np.ndarray:
 69    from skimage.draw import polygon as draw_polygon
 70
 71    labels = np.zeros(shape, dtype="uint16")
 72    instance_id = 0
 73    for annotation in polygons:
 74        if label_choice is not None and annotation["className"] != label_choice:
 75            continue
 76
 77        points = np.asarray(annotation["points"], dtype=float).reshape(-1, 2)
 78        rows, columns = draw_polygon(points[:, 1], points[:, 0], shape=shape)
 79        instance_id += 1
 80        labels[rows, columns] = instance_id
 81    return labels
 82
 83
 84def _process_shard(path: Union[os.PathLike, str], arrow_path: Union[os.PathLike, str]) -> None:
 85    import imageio.v3 as imageio
 86    import pyarrow.ipc as ipc
 87    from PIL import Image
 88    from tqdm import tqdm
 89
 90    with open(arrow_path, "rb") as file:
 91        reader = ipc.open_stream(file)
 92        for batch in tqdm(reader, desc=f"Process {Path(arrow_path).name}"):
 93            columns = batch.select(["image_id", "image", "dataset", "polygon"])
 94            for row in columns.to_pylist():
 95                source = _get_source(row)
 96                if source is None:
 97                    continue
 98
 99                image = np.asarray(Image.open(BytesIO(row["image"]["bytes"])).convert("RGB"))
100                sample_id = _get_sample_id(row, source)
101                image_path = os.path.join(path, "images", source, f"{sample_id}.tif")
102                os.makedirs(os.path.dirname(image_path), exist_ok=True)
103                if not os.path.exists(image_path):
104                    imageio.imwrite(image_path, image, compression="zlib")
105
106                choices = ("cell", "nucleus") if source == "burgess" else ("instances",)
107                for choice in choices:
108                    label_path = os.path.join(path, "labels", source, choice, f"{sample_id}.tif")
109                    if os.path.exists(label_path):
110                        continue
111
112                    os.makedirs(os.path.dirname(label_path), exist_ok=True)
113                    label_choice = choice if source == "burgess" else None
114                    labels = _rasterize(row["polygon"], image.shape[:2], label_choice)
115                    if labels.max() == 0:
116                        raise RuntimeError(f"No '{choice}' polygons found for Micro-Bench sample {sample_id}.")
117                    imageio.imwrite(label_path, labels, compression="zlib")
118
119
120def get_micro_bench_data(path: Union[os.PathLike, str], download: bool = False) -> str:
121    """Download and prepare the segmentation subset of Micro-Bench.
122
123    The seven official test shards contain 3.57 GB in total. Each shard is removed after its
124    segmentation samples have been extracted, so at most one source shard is stored alongside
125    the prepared data.
126
127    Args:
128        path: Filepath to a folder where the downloaded data will be saved.
129        download: Whether to download the data if it is not present.
130
131    Returns:
132        The filepath to the folder with the prepared data.
133    """
134    marker_dir = os.path.join(path, ".processed")
135    os.makedirs(marker_dir, exist_ok=True)
136
137    for filename, checksum in FILES.items():
138        marker_path = os.path.join(marker_dir, filename)
139        if os.path.exists(marker_path):
140            continue
141
142        arrow_path = os.path.join(path, filename)
143        util.download_source(arrow_path, URL.format(filename=filename), download, checksum)
144        _process_shard(path, arrow_path)
145        Path(marker_path).touch()
146        os.remove(arrow_path)
147
148    return str(path)
149
150
151def _validate_source(source: str) -> Tuple[str, str]:
152    if source not in SOURCES:
153        raise ValueError(f"'{source}' is not a valid source. Choose from {list(SOURCES)}.")
154    return SOURCES[source]
155
156
157def get_micro_bench_paths(
158    path: Union[os.PathLike, str],
159    source: Literal[
160        "synthetic_cells",
161        "synthetic_nuclei",
162        "fluo_nuclei",
163        "protein_localization_nuclei",
164        "em_mitochondria",
165    ] = "synthetic_cells",
166    download: bool = False,
167) -> Tuple[List[str], List[str]]:
168    """Get paths to one segmentation source in Micro-Bench.
169
170    Args:
171        path: Filepath to a folder where the downloaded data will be saved.
172        source: The source named by imaging modality and segmentation task. Choose from
173            'synthetic_cells', 'synthetic_nuclei', 'fluo_nuclei', 'protein_localization_nuclei', or
174            'em_mitochondria'. These correspond to Burgess et al., Burgess et al., CellCognition,
175            OpenCell, and Wu et al., respectively.
176        download: Whether to download the data if it is not present.
177
178    Returns:
179        List of filepaths for the image data.
180        List of filepaths for the instance label data.
181    """
182    from natsort import natsorted
183
184    dataset_source, label_choice = _validate_source(source)
185    get_micro_bench_data(path, download)
186
187    image_paths = natsorted(glob(os.path.join(path, "images", dataset_source, "*.tif")))
188    label_paths = natsorted(glob(os.path.join(path, "labels", dataset_source, label_choice, "*.tif")))
189    if not image_paths or len(image_paths) != len(label_paths):
190        raise RuntimeError(f"Could not find matching Micro-Bench images and labels for source '{source}' in {path}.")
191    return image_paths, label_paths
192
193
194def get_micro_bench_dataset(
195    path: Union[os.PathLike, str],
196    patch_shape: Tuple[int, int],
197    source: Literal[
198        "synthetic_cells",
199        "synthetic_nuclei",
200        "fluo_nuclei",
201        "protein_localization_nuclei",
202        "em_mitochondria",
203    ] = "synthetic_cells",
204    download: bool = False,
205    **kwargs,
206) -> Dataset:
207    """Get a Micro-Bench instance segmentation dataset.
208
209    Args:
210        path: Filepath to a folder where the downloaded data will be saved.
211        patch_shape: The patch shape to use for training.
212        source: The source named by imaging modality and segmentation task. See ``get_micro_bench_paths``.
213        download: Whether to download the data if it is not present.
214        kwargs: Additional keyword arguments for ``torch_em.default_segmentation_dataset``.
215
216    Returns:
217        The segmentation dataset.
218    """
219    image_paths, label_paths = get_micro_bench_paths(path, source, download)
220    kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
221    kwargs = util.update_kwargs(kwargs, "ndim", 2)
222
223    return torch_em.default_segmentation_dataset(
224        raw_paths=image_paths,
225        raw_key=None,
226        label_paths=label_paths,
227        label_key=None,
228        patch_shape=patch_shape,
229        is_seg_dataset=False,
230        **kwargs,
231    )
232
233
234def get_micro_bench_loader(
235    path: Union[os.PathLike, str],
236    batch_size: int,
237    patch_shape: Tuple[int, int],
238    source: Literal[
239        "synthetic_cells",
240        "synthetic_nuclei",
241        "fluo_nuclei",
242        "protein_localization_nuclei",
243        "em_mitochondria",
244    ] = "synthetic_cells",
245    download: bool = False,
246    **kwargs,
247) -> DataLoader:
248    """Get a Micro-Bench instance segmentation data loader.
249
250    Args:
251        path: Filepath to a folder where the downloaded data will be saved.
252        batch_size: The batch size for training.
253        patch_shape: The patch shape to use for training.
254        source: The source named by imaging modality and segmentation task. See ``get_micro_bench_paths``.
255        download: Whether to download the data if it is not present.
256        kwargs: Additional keyword arguments for the dataset or PyTorch DataLoader.
257
258    Returns:
259        The DataLoader.
260    """
261    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
262    dataset = get_micro_bench_dataset(
263        path, patch_shape, source, download, **ds_kwargs,
264    )
265    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)