torch_em.data.datasets.electron_microscopy.janelia_nucleus

  1"""The Janelia CellMap nucleus dataset provides nucleus instance segmentation masks for FIB-SEM
  2volumes from the public `janelia-cosem-datasets` S3 bucket (the same "OpenOrganelle" bucket used
  3by other CellMap data), covering several mouse tissues and one Drosophila tissue.
  4
  5For 13 of the 16 volumes, nuclei were segmented automatically with Cellpose 3.0.9 and only
  6cursorily corrected by hand, following the protocol
  7"Generating Nuclei Segmentations for vEM datasets using Cellpose"
  8(https://dx.doi.org/10.17504/protocols.io.n2bvjnq8pgk5/v1). Treat these labels the same way as
  9the MitoNet auto-labels in `mitonet_predicted_kidney.py`: useful pseudo-labels, not verified
 10ground truth.
 11
 12The 3 `jrc_mus-nacc-*` volumes (mouse nucleus accumbens) are instead densely, manually segmented
 13in Amira-Avizo, following
 14"Using Amira to manually segment organelles in vEM for machine learning V.3"
 15(https://dx.doi.org/10.17504/protocols.io.bp2l61rb5vqe/v3). These can be treated as reliable
 16ground truth.
 17
 18Each volume is released as its own Figshare deposit (e.g. https://doi.org/10.6084/m9.figshare.26506555),
 19under CC-BY-4.0. The Figshare entries are metadata-only pointers: their one "file" is a link-only
 20stub whose actual content lives on the S3 bucket, at
 21`s3://janelia-cosem-datasets/{dataset}/{dataset}.zarr/{recon}/em/...` (raw) and
 22`.../{recon}/labels/inference/segmentations/nuc/...` (nucleus instance labels), where `recon` is
 23either "recon-1" or "recon-2" depending on the deposit.
 24The zarr container of `jrc_mus-pancreas-1` has no 'em' group, so its raw data is read from the n5 container
 25`jrc_mus-pancreas-1.n5/em/fibsem-uint16` instead (see `N5_RAW`).
 26
 27The raw EM data is stored as a multiscale OME-Zarr pyramid at up to 4 nm isotropic resolution,
 28while the nucleus labels are only released at one, coarser resolution (their pyramid level "s0").
 29This module finds the raw pyramid level whose resolution matches the label's "s0" level (their
 30shapes then match exactly, so no rescaling is needed, unlike the mito/tissue labels in
 31`hemibrain.py` which need explicit upsampling) and downloads only that matched pair, in full,
 32for each requested dataset name.
 33
 34Please cite the Cellpose or Amira-Avizo protocol above (as appropriate for the chosen dataset)
 35and the CellMap project (https://www.janelia.org/project-team/cellmap) if you use this data.
 36"""
 37
 38import os
 39import re
 40from typing import List, Optional, Tuple, Union
 41
 42import numpy as np
 43
 44from torch.utils.data import DataLoader, Dataset
 45
 46import torch_em
 47
 48from .. import util
 49
 50
 51BUCKET_URL = "https://janelia-cosem-datasets.s3.amazonaws.com/"
 52
 53# dataset name -> (figshare article id, recon folder, annotation type)
 54# "cellpose" = automatic Cellpose 3.0.9 prediction with cursory manual correction.
 55# "manual" = dense manual segmentation in Amira-Avizo.
 56DATASETS = {
 57    "jrc_mus-pancreas-1": (26506552, "recon-1", "cellpose"),
 58    "jrc_mus-pancreas-2": (26506555, "recon-1", "cellpose"),
 59    "jrc_mus-pancreas-3": (26506561, "recon-1", "cellpose"),
 60    "jrc_mus-granule-neurons-1": (26506510, "recon-2", "cellpose"),
 61    "jrc_mus-granule-neurons-2": (26506513, "recon-2", "cellpose"),
 62    "jrc_mus-granule-neurons-3": (26506516, "recon-2", "cellpose"),
 63    "jrc_mus-guard-hair-follicle": (26506519, "recon-1", "cellpose"),
 64    "jrc_mus-kidney": (26506522, "recon-1", "cellpose"),
 65    "jrc_mus-kidney-2": (26506534, "recon-1", "cellpose"),
 66    "jrc_mus-liver-2": (26506537, "recon-1", "cellpose"),
 67    "jrc_mus-meissner-corpuscle-1": (26506540, "recon-1", "cellpose"),
 68    "jrc_mus-meissner-corpuscle-2": (26506543, "recon-1", "cellpose"),
 69    "jrc_fly-mb-z0419-20": (26506507, "recon-1", "cellpose"),
 70    "jrc_mus-nacc-2": (26513209, "recon-2", "manual"),
 71    "jrc_mus-nacc-3": (26513212, "recon-2", "manual"),
 72    "jrc_mus-nacc-4": (26513215, "recon-2", "manual"),
 73}
 74
 75# dataset name -> raw array in the n5 container, for datasets without an 'em' group in the zarr container.
 76N5_RAW = {"jrc_mus-pancreas-1": "em/fibsem-uint16"}
 77
 78
 79def _get_json(url):
 80    import requests
 81    resp = requests.get(url, timeout=60)
 82    resp.raise_for_status()
 83    return resp.json()
 84
 85
 86def _list_prefixes(prefix):
 87    import requests
 88    resp = requests.get(BUCKET_URL, params={"prefix": prefix, "delimiter": "/"}, timeout=60)
 89    resp.raise_for_status()
 90    return re.findall(r"<Prefix>([^<]*)</Prefix>", resp.text)
 91
 92
 93def _find_em_array_name(dataset_name, recon):
 94    prefix = f"{dataset_name}/{dataset_name}.zarr/{recon}/em/"
 95    names = [p[len(prefix):].rstrip("/") for p in _list_prefixes(prefix) if p != prefix]
 96    if not names:
 97        raise RuntimeError(f"Could not find an 'em' array for '{dataset_name}' under '{prefix}'.")
 98    return names[0]
 99
100
101def _get_nucleus_scale(dataset_name, recon):
102    nuc_attrs = _get_json(
103        f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/labels/inference/segmentations/nuc/.zattrs"
104    )
105    return nuc_attrs["multiscales"][0]["datasets"][0]["coordinateTransformations"][0]["scale"]
106
107
108def _find_matching_em_level(dataset_name, recon, em_name):
109    em_attrs = _get_json(f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}/.zattrs")
110    nuc_scale = _get_nucleus_scale(dataset_name, recon)
111    for level in em_attrs["multiscales"][0]["datasets"]:
112        scale = level["coordinateTransformations"][0]["scale"]
113        if scale == nuc_scale:
114            return level["path"]
115    raise RuntimeError(
116        f"No EM resolution level of '{dataset_name}' matches the nucleus label resolution {nuc_scale}."
117    )
118
119
120def _read_n5_raw(dataset_name, recon, em_path):
121    """Read the raw data at the resolution of the nucleus labels from the n5 container (gzip blocks, big-endian)."""
122    import gzip
123    import struct
124    import requests
125
126    n5_path = f"{dataset_name}/{dataset_name}.n5/{em_path}"
127    nuc_scale = _get_nucleus_scale(dataset_name, recon)
128    levels = _get_json(f"{BUCKET_URL}{n5_path}/attributes.json")["multiscales"][0]["datasets"]
129    # The scales of the n5 levels differ from the label scale by float rounding.
130    matching = [level["path"] for level in levels if np.allclose(level["transform"]["scale"], nuc_scale)]
131    if not matching:
132        raise RuntimeError(f"No n5 level of '{dataset_name}' matches the nucleus label resolution {nuc_scale}.")
133    level_url = f"{BUCKET_URL}{n5_path}/{matching[0]}"
134    attrs = _get_json(f"{level_url}/attributes.json")
135    assert attrs["compression"]["type"] == "gzip", attrs["compression"]
136
137    # The n5 axis order is (x, y, z), so the shapes and the block ids are reversed.
138    shape, block_shape = attrs["dimensions"][::-1], attrs["blockSize"][::-1]
139    dtype = np.dtype(attrs["dataType"])
140    raw = np.zeros(shape, dtype=dtype)
141    for block_id in np.ndindex(*[int(np.ceil(s / b)) for s, b in zip(shape, block_shape)]):
142        resp = requests.get(f"{level_url}/" + "/".join(str(i) for i in block_id[::-1]), timeout=120)
143        if resp.status_code == 404:  # A missing block is empty.
144            continue
145        resp.raise_for_status()
146        mode, ndim = struct.unpack(">HH", resp.content[:4])
147        block_dims = struct.unpack(f">{ndim}I", resp.content[4:4 + 4 * ndim])[::-1]
148        data_start = 4 + 4 * ndim + (4 if mode == 1 else 0)
149        data = gzip.decompress(resp.content[data_start:])
150        block = np.frombuffer(data, dtype=dtype.newbyteorder(">")).reshape(block_dims)
151        bb = tuple(slice(i * b, min(i * b + d, s)) for i, b, d, s in zip(block_id, block_shape, block_dims, shape))
152        raw[bb] = block[tuple(slice(0, b.stop - b.start) for b in bb)]
153
154    return raw, f"s3://janelia-cosem-datasets/{n5_path}/{matching[0]}"
155
156
157def _read_zarr_array(s3_path):
158    import zarr
159    import fsspec
160
161    store = fsspec.get_mapper(s3_path, anon=True)
162    array = zarr.open(store, mode="r")
163    return np.asarray(array[:])
164
165
166def get_janelia_nucleus_data(path: Union[os.PathLike, str], dataset_name: str, download: bool = False) -> str:
167    """Download nucleus instance segmentation data for one Janelia CellMap dataset.
168
169    Args:
170        path: Filepath to a folder where the cached zarr store will be saved.
171        dataset_name: The name of the dataset, one of the keys in `DATASETS`.
172        download: Whether to download the data if it is not present.
173
174    Returns:
175        The filepath to the cached zarr store.
176    """
177    import zarr
178    from zarr.codecs import BloscCodec
179
180    if dataset_name not in DATASETS:
181        raise ValueError(f"'{dataset_name}' is not a valid dataset name. Choose from {sorted(DATASETS.keys())}.")
182
183    os.makedirs(str(path), exist_ok=True)
184    zarr_path = os.path.join(str(path), f"{dataset_name}.zarr")
185
186    root = zarr.open_group(zarr_path, mode="a")
187    if "raw" in root and "labels" in root:
188        return zarr_path
189
190    if not download:
191        raise RuntimeError(f"No cached data found at '{zarr_path}'. Set download=True to stream it from S3.")
192
193    _, recon, annotation = DATASETS[dataset_name]
194    nuc_s3 = (
195        f"s3://janelia-cosem-datasets/{dataset_name}/{dataset_name}.zarr/{recon}/"
196        "labels/inference/segmentations/nuc/s0"
197    )
198
199    print(f"Streaming '{dataset_name}' ({recon}, {annotation} labels) from the janelia-cosem-datasets S3 bucket ...")
200    if dataset_name in N5_RAW:
201        raw, em_s3 = _read_n5_raw(dataset_name, recon, N5_RAW[dataset_name])
202    else:
203        em_name = _find_em_array_name(dataset_name, recon)
204        em_level = _find_matching_em_level(dataset_name, recon, em_name)
205        em_s3 = f"s3://janelia-cosem-datasets/{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}/{em_level}"
206        raw = _read_zarr_array(em_s3)
207    labels = _read_zarr_array(nuc_s3)
208    assert raw.shape == labels.shape, f"Shape mismatch for '{dataset_name}': raw {raw.shape} vs labels {labels.shape}"
209
210    def _make_array(name, data, shuffle):
211        arr = root.create_array(
212            name, shape=data.shape, chunks=tuple(min(128, s) for s in data.shape), dtype=data.dtype,
213            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
214        )
215        arr[:] = data
216
217    root.attrs["dataset"] = dataset_name
218    root.attrs["recon"] = recon
219    root.attrs["annotation_type"] = annotation
220    root.attrs["labels_are_manual"] = annotation == "manual"
221    root.attrs["em_source"] = em_s3
222    root.attrs["label_source"] = nuc_s3
223
224    _make_array("raw", raw, shuffle="shuffle")
225    _make_array("labels", labels, shuffle="bitshuffle")
226
227    print(f"Cached '{dataset_name}' to '{zarr_path}' (shape {raw.shape}).")
228    return zarr_path
229
230
231def get_janelia_nucleus_paths(
232    path: Union[os.PathLike, str], dataset_names: Optional[List[str]] = None, download: bool = False,
233) -> List[str]:
234    """Get paths to cached Janelia CellMap nucleus zarr stores.
235
236    Args:
237        path: Filepath to a folder where the cached zarr stores will be saved.
238        dataset_names: The names of the datasets to use. Defaults to all datasets in `DATASETS`.
239        download: Whether to download the data if it is not present.
240
241    Returns:
242        List of filepaths to the cached zarr stores.
243    """
244    if dataset_names is None:
245        dataset_names = list(DATASETS.keys())
246    return [get_janelia_nucleus_data(path, name, download) for name in dataset_names]
247
248
249def get_janelia_nucleus_dataset(
250    path: Union[os.PathLike, str],
251    patch_shape: Tuple[int, int, int],
252    dataset_names: Optional[List[str]] = None,
253    download: bool = False,
254    offsets: Optional[List[List[int]]] = None,
255    boundaries: bool = False,
256    **kwargs,
257) -> Dataset:
258    """Get the Janelia CellMap nucleus dataset for nucleus instance segmentation.
259
260    Args:
261        path: Filepath to a folder where the cached zarr stores will be saved.
262        patch_shape: The patch shape (z, y, x) to use for training.
263        dataset_names: The names of the datasets to use. Defaults to all datasets in `DATASETS`.
264        download: Whether to download the data if it is not present.
265        offsets: Offset values for affinity computation used as target.
266        boundaries: Whether to compute boundaries as the target.
267        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
268
269    Returns:
270        The segmentation dataset.
271    """
272    assert len(patch_shape) == 3
273
274    paths = get_janelia_nucleus_paths(path, dataset_names, download)
275
276    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
277    kwargs, _ = util.add_instance_label_transform(
278        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
279    )
280
281    return torch_em.default_segmentation_dataset(
282        raw_paths=paths,
283        raw_key="raw",
284        label_paths=paths,
285        label_key="labels",
286        patch_shape=patch_shape,
287        **kwargs,
288    )
289
290
291def get_janelia_nucleus_loader(
292    path: Union[os.PathLike, str],
293    patch_shape: Tuple[int, int, int],
294    batch_size: int,
295    dataset_names: Optional[List[str]] = None,
296    download: bool = False,
297    offsets: Optional[List[List[int]]] = None,
298    boundaries: bool = False,
299    **kwargs,
300) -> DataLoader:
301    """Get the DataLoader for nucleus instance segmentation in the Janelia CellMap nucleus dataset.
302
303    Args:
304        path: Filepath to a folder where the cached zarr stores will be saved.
305        patch_shape: The patch shape (z, y, x) to use for training.
306        batch_size: The batch size for training.
307        dataset_names: The names of the datasets to use. Defaults to all datasets in `DATASETS`.
308        download: Whether to download the data if it is not present.
309        offsets: Offset values for affinity computation used as target.
310        boundaries: Whether to compute boundaries as the target.
311        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
312            or for the PyTorch DataLoader.
313
314    Returns:
315        The DataLoader.
316    """
317    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
318    dataset = get_janelia_nucleus_dataset(
319        path, patch_shape, dataset_names=dataset_names, download=download,
320        offsets=offsets, boundaries=boundaries, **ds_kwargs
321    )
322    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)