torch_em.data.datasets.electron_microscopy.mitonet_predicted_kidney

  1"""The MitoNet-Predicted-Kidney dataset streams mitochondria instance segmentation for the Janelia
  2OpenOrganelle mouse kidney FIB-SEM volume (jrc_mus-kidney) at 16 nm isotropic resolution.
  3
  4IMPORTANT: The labels are automatically generated instance segmentation predictions from MitoNet
  5(Conrad and Narayan 2022), not manually annotated ground truth. They are hardened at a semantic
  6confidence threshold of 0.5 and a contour confidence threshold of 0.3, followed by connected components
  7and watershed. Use them as weak or pseudo-labels, not as verified ground truth.
  8
  9Raw data is streamed from the public OpenOrganelle S3 bucket. Predicted labels are streamed from a
 10remote Zarr array embedded in a Zenodo/figshare ZIP archive, decoded chunk-by-chunk via HTTP range
 11requests. Only the requested bounding box is downloaded and cached locally as a zarr v3 store.
 12
 13The predictions are available at https://doi.org/10.6084/m9.figshare.20749729, licensed CC-BY-4.0.
 14The MitoNet method is described in https://doi.org/10.1016/j.cels.2022.12.004.
 15Please cite this publication if you use the predictions in your research.
 16"""
 17
 18import os
 19from typing import List, Tuple, Union
 20
 21import numpy as np
 22
 23from torch.utils.data import DataLoader, Dataset
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30RAW_S3_URL = "s3://janelia-cosem-datasets/jrc_mus-kidney/jrc_mus-kidney.zarr/recon-1/em/fibsem-uint8/s1"
 31LABEL_ZIP_URL = "https://ndownloader.figshare.com/files/36985378"
 32LABEL_ARRAY_PATH = "kidney16nm.zarr/empanada_mito_pred"
 33
 34FULL_SHAPE = (11099, 3988, 6143)  # (z, y, x) at 16 nm isotropic resolution.
 35LABEL_CHUNK_SHAPE = (512, 512, 512)
 36
 37
 38def _bbox_hash(bounding_box):
 39    import hashlib
 40    return hashlib.md5("_".join(str(v) for v in bounding_box).encode()).hexdigest()[:12]
 41
 42
 43def _read_raw_block(z_min, z_max, y_min, y_max, x_min, x_max):
 44    """Slice the requested block directly from the public OpenOrganelle S3 zarr array."""
 45    import zarr
 46    import fsspec
 47
 48    store = fsspec.get_mapper(RAW_S3_URL, anon=True)
 49    raw = zarr.open(store, mode="r")
 50    return np.asarray(raw[z_min:z_max, y_min:y_max, x_min:x_max])
 51
 52
 53def _read_label_block(z_min, z_max, y_min, y_max, x_min, x_max):
 54    """Decode only the chunks overlapping the requested block from the remote ZIP-embedded zarr array.
 55
 56    `zarr.open` cannot be used directly here: the archive's zarr v2 store has no `.zattrs` entry, and
 57    `fsspec`'s zip backend raises instead of treating that as optional. The array layout (shape, chunk
 58    shape, dtype, blosc codec) is fixed and taken from the published `.zarray` metadata instead.
 59    """
 60    import zipfile
 61    import fsspec
 62    import numcodecs
 63
 64    fs = fsspec.filesystem("http")
 65    zf = zipfile.ZipFile(fs.open(LABEL_ZIP_URL, "rb"))
 66    codec = numcodecs.Blosc()
 67
 68    cz, cy, cx = LABEL_CHUNK_SHAPE
 69    out = np.zeros((z_max - z_min, y_max - y_min, x_max - x_min), dtype="<u4")
 70
 71    for iz in range(z_min // cz, -(-z_max // cz)):
 72        for iy in range(y_min // cy, -(-y_max // cy)):
 73            for ix in range(x_min // cx, -(-x_max // cx)):
 74                name = f"{LABEL_ARRAY_PATH}/{iz}.{iy}.{ix}"
 75                if name not in zf.namelist():
 76                    continue
 77                chunk = np.frombuffer(codec.decode(zf.read(name)), dtype="<u4").reshape(LABEL_CHUNK_SHAPE)
 78
 79                cz0, cy0, cx0 = iz * cz, iy * cy, ix * cx
 80                sz = slice(max(z_min, cz0) - cz0, min(z_max, cz0 + cz) - cz0)
 81                sy = slice(max(y_min, cy0) - cy0, min(y_max, cy0 + cy) - cy0)
 82                sx = slice(max(x_min, cx0) - cx0, min(x_max, cx0 + cx) - cx0)
 83                oz = slice(max(z_min, cz0) - z_min, min(z_max, cz0 + cz) - z_min)
 84                oy = slice(max(y_min, cy0) - y_min, min(y_max, cy0 + cy) - y_min)
 85                ox = slice(max(x_min, cx0) - x_min, min(x_max, cx0 + cx) - x_min)
 86                out[oz, oy, ox] = chunk[sz, sy, sx]
 87
 88    return out
 89
 90
 91def get_mitonet_predicted_kidney_data(
 92    path: Union[os.PathLike, str], bounding_box: Tuple[int, int, int, int, int, int], download: bool = False,
 93) -> str:
 94    """Stream a subvolume of the MitoNet-predicted mouse kidney data and cache it as a zarr v3 store.
 95
 96    Args:
 97        path: Filepath to a folder where the cached zarr store will be saved.
 98        bounding_box: The region to fetch as (z_min, z_max, y_min, y_max, x_min, x_max)
 99            in voxel coordinates at 16 nm isotropic resolution.
100        download: Whether to stream and cache the data if it is not present.
101
102    Returns:
103        The filepath to the cached zarr store.
104    """
105    import zarr
106    from zarr.codecs import BloscCodec
107
108    os.makedirs(str(path), exist_ok=True)
109    zarr_path = os.path.join(str(path), f"kidney_{_bbox_hash(bounding_box)}.zarr")
110
111    root = zarr.open_group(zarr_path, mode="a")
112    if "raw" in root and "labels" in root:
113        return zarr_path
114
115    if not download:
116        raise RuntimeError(f"No cached data found at '{zarr_path}'. Set download=True to stream it.")
117
118    z_min, z_max, y_min, y_max, x_min, x_max = bounding_box
119    for bound, limit in zip((z_max, y_max, x_max), FULL_SHAPE):
120        assert bound <= limit, f"Bounding box {bounding_box} exceeds the full volume shape {FULL_SHAPE}"
121
122    raw_block = _read_raw_block(z_min, z_max, y_min, y_max, x_min, x_max)
123    label_block = _read_label_block(z_min, z_max, y_min, y_max, x_min, x_max)
124    assert raw_block.shape == label_block.shape, f"Shape mismatch: {raw_block.shape} vs {label_block.shape}"
125
126    def _make_array(name, data, shuffle):
127        arr = root.create_array(
128            name, shape=data.shape, chunks=(64, 256, 256), dtype=data.dtype,
129            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
130        )
131        arr[:] = data
132
133    root.attrs["bounding_box"] = list(bounding_box)
134    root.attrs["resolution_nm"] = [16, 16, 16]
135    root.attrs["labels_are_automatic_predictions"] = True
136
137    _make_array("raw", raw_block, shuffle="shuffle")
138    _make_array("labels", label_block, shuffle="bitshuffle")
139
140    return zarr_path
141
142
143def get_mitonet_predicted_kidney_paths(
144    path: Union[os.PathLike, str],
145    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
146    download: bool = False,
147) -> List[str]:
148    """Get paths to cached MitoNet-predicted mouse kidney zarr stores.
149
150    Args:
151        path: Filepath to a folder where the cached zarr stores will be saved.
152        bounding_boxes: List of regions to fetch, each as
153            (z_min, z_max, y_min, y_max, x_min, x_max) in voxel coordinates at 16 nm resolution.
154        download: Whether to stream and cache the data if it is not present.
155
156    Returns:
157        List of filepaths to the cached zarr stores.
158    """
159    return [get_mitonet_predicted_kidney_data(path, bbox, download) for bbox in bounding_boxes]
160
161
162def get_mitonet_predicted_kidney_dataset(
163    path: Union[os.PathLike, str],
164    patch_shape: Tuple[int, int, int],
165    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
166    download: bool = False,
167    **kwargs,
168) -> Dataset:
169    """Get the MitoNet-predicted mouse kidney dataset for mitochondria instance segmentation.
170
171    The labels are automatically generated MitoNet predictions, not manual ground truth.
172
173    Args:
174        path: Filepath to a folder where the cached zarr stores will be saved.
175        patch_shape: The patch shape (z, y, x) to use for training.
176        bounding_boxes: List of subvolumes to use, each as
177            (z_min, z_max, y_min, y_max, x_min, x_max) in 16 nm voxel coordinates.
178        download: Whether to stream and cache data if not already present.
179        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
180
181    Returns:
182        The segmentation dataset.
183    """
184    assert len(patch_shape) == 3
185
186    paths = get_mitonet_predicted_kidney_paths(path, bounding_boxes, download)
187    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
188
189    return torch_em.default_segmentation_dataset(
190        raw_paths=paths,
191        raw_key="raw",
192        label_paths=paths,
193        label_key="labels",
194        patch_shape=patch_shape,
195        **kwargs,
196    )
197
198
199def get_mitonet_predicted_kidney_loader(
200    path: Union[os.PathLike, str],
201    patch_shape: Tuple[int, int, int],
202    batch_size: int,
203    bounding_boxes: List[Tuple[int, int, int, int, int, int]],
204    download: bool = False,
205    **kwargs,
206) -> DataLoader:
207    """Get the DataLoader for mitochondria instance segmentation in the MitoNet-predicted mouse kidney data.
208
209    The labels are automatically generated MitoNet predictions, not manual ground truth.
210
211    Args:
212        path: Filepath to a folder where the cached zarr stores will be saved.
213        patch_shape: The patch shape (z, y, x) to use for training.
214        batch_size: The batch size for training.
215        bounding_boxes: List of subvolumes to use, each as
216            (z_min, z_max, y_min, y_max, x_min, x_max) in 16 nm voxel coordinates.
217        download: Whether to stream and cache data if not already present.
218        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
219            or for the PyTorch DataLoader.
220
221    Returns:
222        The DataLoader.
223    """
224    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
225    dataset = get_mitonet_predicted_kidney_dataset(path, patch_shape, bounding_boxes, download, **ds_kwargs)
226    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)