torch_em.data.datasets.electron_microscopy.openorganelle_nucleus

  1"""This dataset provides nucleus segmentation masks for 30 volumes from the public
  2"OpenOrganelle" `janelia-cosem-datasets` S3 bucket (the same bucket used by `cellmap.py` and
  3`janelia_nucleus.py`), none of which are covered by either of those existing loaders. All 30 are
  4released under CC0-1.0, except `jrc_mus-thymus-1` and `aic_desmosome-3` whose license and paper
  5could not be independently verified (`aic_desmosome-3` is an Allen Institute dataset, so the
  6Janelia CC0-1.0 convention documented for the rest does not apply to it).
  7
  8Three label folder naming conventions are used across the bucket, all distinct from
  9`janelia_nucleus.py`'s:
 10
 11- `labels/inference/nucleus_seg/{level}` (`jrc_ctl-id8-2`, `jrc_choroid-plexus-2`, `jrc_dauer-larva`).
 12- `labels/inference/segmentations/nuc/{level}` (most of the remaining 27 datasets).
 13- `labels/inference/segmentations/nucleus/{level}` (`jrc_hum-airway-14953vc` only).
 14- `labels/inference/segmentations/nuc_mem/{level}` (`jrc_mus-liver-4/5/6`) - these three are
 15  nuclear MEMBRANE labels, a distinct annotation target from a nucleus mask/instance segmentation,
 16  tracked via `DATASETS[name]["label_target"]`.
 17
 18Label and raw data are usually released at an exactly matching multiscale pyramid level (no
 19resampling, same axis order). Three exceptions are handled: `jrc_dauer-larva` and
 20`jrc_hela-h89-1/2` have no exact-scale match and are resampled onto the label grid via
 21`scipy.ndimage.zoom`; `jrc_choroid-plexus-2` and `jrc_mus-epididymis-1/2` store one of the two
 22arrays with a reversed axis order, corrected via `np.transpose`.
 23
 24Datasets whose full matched-resolution raw+label pair is small (roughly under 5 GB) are
 25downloaded whole, following `janelia_nucleus.py`'s pattern. The rest (full volumes ranging from
 26tens of GB to several TB) are downloaded as one fixed-size representative crop, auto-detected
 27from a non-empty region of the coarsest available label pyramid level, following `cellmap.py`'s
 28per-crop pattern - see `DATASETS[name]["full_array"]`.
 29Alternatively, the full raw and label arrays of a dataset are downloaded at one pyramid level, selected by its voxel
 30size in nm (`voxel_size`). They are cached as `<name>_<voxel_size>nm.zarr` and cropped to their common shape.
 31
 32`jrc_mosquito-stylet-6` is the only dataset whose raw and label arrays are stored in zarr v3 with
 33sharding (1024^3 outer shards, 64^3 inner chunks); `_open_remote_zarr` uses `zarr.storage.
 34FsspecStore` rather than `fsspec.get_mapper` specifically so that a crop only pulls the individual
 35inner chunks it needs via ranged reads, not whole ~1 GB shard files.
 36
 37Please cite the CellMap project (https://www.janelia.org/project-team/cellmap) if you use this
 38data (except `aic_desmosome-3`, an Allen Institute for Cell Science dataset).
 39"""
 40
 41import os
 42from typing import List, Optional, Tuple, Union
 43
 44import numpy as np
 45
 46from torch.utils.data import DataLoader, Dataset
 47
 48import torch_em
 49
 50from .. import util
 51
 52
 53BUCKET_URL = "https://janelia-cosem-datasets.s3.amazonaws.com/"
 54
 55# dataset name -> recon folder, em array name, label subpath (relative to "labels/inference/"),
 56# label pyramid level, label kind, label target, axis to reverse ("none", "label", or "raw"),
 57# and whether the full matched-resolution pair is downloaded whole (True) or as one fixed-size
 58# crop auto-detected from a non-empty region (False).
 59# label_kind: "binary" = label array only ever takes values {0, 1}; "instance" = per-nucleus IDs.
 60# label_target: "nucleus" (a nucleus mask/instance segmentation) or "nuclear_membrane" (a nuclear
 61# envelope surface label, a distinct annotation target).
 62DATASETS = {
 63    "jrc_ctl-id8-2": {
 64        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "nucleus_seg", "label_level": "s0",
 65        "label_kind": "binary", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 66        "bounding_box": ((3824, 4048), (448, 672), (4064, 4288)),
 67    },
 68    "jrc_choroid-plexus-2": {
 69        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "nucleus_seg", "label_level": "s1",
 70        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "label", "full_array": False,
 71        "bounding_box": ((780, 1100), (2220, 2540), (620, 940)),
 72    },
 73    "jrc_dauer-larva": {
 74        "recon": "recon-1", "em_name": "tem-uint8", "label_subpath": "nucleus_seg", "label_level": "s0",
 75        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 76        "bounding_box": ((130, 330), (1804, 2004), (11724, 11924)),
 77    },
 78    "jrc_ccl81-covid-1": {
 79        "recon": "recon-1", "em_name": "fibsem-uint16", "label_subpath": "segmentations/nuc", "label_level": "s0",
 80        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 81    },
 82    "jrc_cos7-11": {
 83        "recon": "recon-1", "em_name": "fibsem-uint16", "label_subpath": "segmentations/nuc", "label_level": "s0",
 84        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 85    },
 86    "jrc_ctl-id8-3": {
 87        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
 88        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 89    },
 90    "jrc_ctl-id8-4": {
 91        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
 92        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 93    },
 94    "jrc_ctl-id8-5": {
 95        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
 96        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
 97    },
 98    "jrc_fly-acc-calyx-1": {
 99        "recon": "recon-1", "em_name": "fibsem-uint16", "label_subpath": "segmentations/nuc", "label_level": "s0",
100        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
101    },
102    "jrc_fly-fsb-1": {
103        "recon": "recon-1", "em_name": "fibsem-uint16", "label_subpath": "segmentations/nuc", "label_level": "s0",
104        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
105    },
106    "jrc_hela-21": {
107        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
108        "label_kind": "binary", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
109    },
110    "jrc_hela-22": {
111        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
112        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
113    },
114    "jrc_hela-bfa": {
115        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
116        "label_kind": "binary", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
117    },
118    "jrc_hela-h89-1": {
119        "recon": "recon-1", "em_name": "fibsem-uint16", "label_subpath": "segmentations/nuc", "label_level": "s0",
120        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
121    },
122    "jrc_hela-h89-2": {
123        "recon": "recon-1", "em_name": "fibsem-uint16", "label_subpath": "segmentations/nuc", "label_level": "s0",
124        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
125    },
126    "jrc_hela-nz-1": {
127        "recon": "recon-2", "em_name": "fibsem-int16", "label_subpath": "segmentations/nuc", "label_level": "s0",
128        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
129    },
130    "jrc_hum-airway-14953vc": {
131        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nucleus", "label_level": "s0",
132        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
133    },
134    # Stored in zarr v3 with sharding (1024^3 outer shards, 64^3 inner chunks); `_open_remote_zarr`
135    # uses `FsspecStore` rather than `fsspec.get_mapper` so a crop only pulls the inner chunks it
136    # actually needs, not whole ~1GB shard files. Its label pyramid has only one level (s0, full
137    # resolution) - there is no genuinely coarse level to cheaply scan for a non-empty region, so
138    # this uses a verified fixed crop instead of auto-detection.
139    "jrc_mosquito-stylet-6": {
140        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
141        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
142        "bounding_box": ((1088, 1216), (640, 768), (768, 896)),
143    },
144    "jrc_mus-epididymis-1": {
145        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
146        "label_kind": "binary", "label_target": "nucleus", "axis_reverse": "raw", "full_array": True,
147    },
148    "jrc_mus-epididymis-2": {
149        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
150        "label_kind": "binary", "label_target": "nucleus", "axis_reverse": "raw", "full_array": True,
151    },
152    "jrc_mus-hippocampus-1": {
153        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
154        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
155    },
156    # These three membranes are too thin/sparse to survive downsampling to the coarsest label
157    # pyramid level, so auto-crop-detection finds no non-empty region; use a verified fixed crop.
158    "jrc_mus-liver-4": {
159        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc_mem", "label_level": "s0",
160        "label_kind": "binary", "label_target": "nuclear_membrane", "axis_reverse": "none", "full_array": False,
161        "bounding_box": ((1409, 1537), (2000, 2128), (6000, 6128)),
162    },
163    "jrc_mus-liver-5": {
164        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc_mem", "label_level": "s0",
165        "label_kind": "binary", "label_target": "nuclear_membrane", "axis_reverse": "none", "full_array": False,
166        "bounding_box": ((2961, 3089), (5096, 5224), (4553, 4681)),
167    },
168    "jrc_mus-liver-6": {
169        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc_mem", "label_level": "s0",
170        "label_kind": "binary", "label_target": "nuclear_membrane", "axis_reverse": "none", "full_array": False,
171        "bounding_box": ((5600, 5792), (3970, 4162), (6610, 6802)),
172    },
173    "jrc_mus-pancreas-4": {
174        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
175        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
176    },
177    "jrc_mus-sc-zp104a": {
178        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
179        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
180    },
181    "jrc_mus-sc-zp105a": {
182        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
183        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
184    },
185    "jrc_mus-skin-1": {
186        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
187        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": False,
188    },
189    "jrc_mus-thymus-1": {
190        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
191        "label_kind": "instance", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
192    },
193    "aic_desmosome-3": {
194        "recon": "recon-1", "em_name": "fibsem-uint8", "label_subpath": "segmentations/nuc", "label_level": "s0",
195        "label_kind": "binary", "label_target": "nucleus", "axis_reverse": "none", "full_array": True,
196    },
197}
198
199# Datasets whose data license and paper could not be independently verified. Everything else in
200# `DATASETS` is CC0-1.0 per the OpenOrganelle/CellMap project convention.
201UNVERIFIED_LICENSE = ("jrc_mus-thymus-1", "aic_desmosome-3")
202
203CROP_SHAPE_DEFAULT = (128, 128, 128)
204
205
206def _open_remote_zarr(s3_path):
207    import zarr
208    from zarr.storage import FsspecStore
209
210    # `FsspecStore` (not `fsspec.get_mapper`) is required for `jrc_mosquito-stylet-6`'s zarr v3
211    # sharded array: it performs ranged partial reads of individual inner chunks, whereas
212    # `fsspec.get_mapper` fetches whole shard files (up to ~1GB each) even for a small crop.
213    store = FsspecStore.from_url(s3_path, storage_options={"anon": True}, read_only=True)
214    return zarr.open(store, mode="r")
215
216
217def _get_json(url):
218    import requests
219
220    resp = requests.get(url, timeout=60)
221    resp.raise_for_status()
222    return resp.json()
223
224
225def _multiscale_levels(group_url):
226    """Read multiscale pyramid levels from a zarr group, trying zarr v2 and v3 metadata.
227
228    `jrc_mosquito-stylet-6`'s raw and label arrays are stored in zarr v3 (`zarr.json`), unlike
229    every other dataset here (zarr v2, `.zattrs`); its label metadata nests multiscales under
230    `attributes.ome`, while its raw metadata puts them directly under `attributes`.
231    """
232    try:
233        multiscales = _get_json(f"{group_url}/.zattrs")["multiscales"]
234    except Exception:
235        attributes = _get_json(f"{group_url}/zarr.json")["attributes"]
236        multiscales = attributes["ome"]["multiscales"] if "ome" in attributes else attributes["multiscales"]
237    return {
238        d["path"]: tuple(d["coordinateTransformations"][0]["scale"])
239        for d in multiscales[0]["datasets"]
240    }
241
242
243def _find_matching_em_level(dataset_name, recon, em_name, label_scale):
244    """Find the EM pyramid level whose scale exactly matches the label pyramid level, or `None`."""
245    em_levels = _multiscale_levels(f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}")
246    for level, scale in em_levels.items():
247        if scale == label_scale:
248            return level, scale
249    return None, None
250
251
252def _nearest_em_level(dataset_name, recon, em_name, label_scale):
253    """Find the EM pyramid level whose scale is closest to the label pyramid level.
254
255    Used only for datasets whose label pyramid has no exact scale match.
256    """
257    em_levels = _multiscale_levels(f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}")
258    best_level, best_dist = None, None
259    for level, scale in em_levels.items():
260        dist = sum((s - lv) ** 2 for s, lv in zip(scale, label_scale))
261        if best_dist is None or dist < best_dist:
262            best_level, best_dist = level, dist
263    return best_level, em_levels[best_level]
264
265
266def _find_level(levels, voxel_size, array_path):
267    """Find the pyramid level with the given isotropic voxel size."""
268    for level, scale in levels.items():
269        if np.allclose(scale, voxel_size):
270            return level, scale
271    raise ValueError(f"'{array_path}' has no pyramid level with voxel size {voxel_size} nm. The levels are {levels}.")
272
273
274def _auto_bounding_box(dataset_name, recon, label_full_path, label_level, label_scale, crop_shape):
275    """Center a crop on a non-empty region, found via the coarsest available label level.
276
277    Downloading the coarsest level (rather than `label_level`, which may be a huge full-resolution
278    array) keeps this cheap even for multi-TB datasets.
279    """
280    levels = _multiscale_levels(f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/{label_full_path}")
281    coarsest_level = max(levels, key=lambda lv: levels[lv])
282    coarse_url = (
283        f"s3://janelia-cosem-datasets/{dataset_name}/{dataset_name}.zarr/{recon}/"
284        f"{label_full_path}/{coarsest_level}"
285    )
286    coarse_arr = _open_remote_zarr(coarse_url)
287    coarse_data = np.asarray(coarse_arr[:])
288    nonzero = np.argwhere(coarse_data > 0)
289    if len(nonzero) == 0:
290        raise RuntimeError(f"No non-empty '{label_full_path}' region found for '{dataset_name}'.")
291    center_coarse = nonzero[len(nonzero) // 2]
292    scale_factor = np.array(levels[coarsest_level]) / np.array(label_scale)
293    center = (center_coarse * scale_factor).astype(int)
294    half = [c // 2 for c in crop_shape]
295    return tuple((max(0, int(c - h)), int(c + h)) for c, h in zip(center, half))
296
297
298def get_openorganelle_nucleus_data(
299    path: Union[os.PathLike, str], dataset_name: str, download: bool = False, voxel_size: Optional[float] = None,
300) -> str:
301    """Download nucleus segmentation data for one OpenOrganelle dataset.
302
303    Args:
304        path: Filepath to a folder where the cached zarr store will be saved.
305        dataset_name: The name of the dataset, one of the keys in `DATASETS`.
306        download: Whether to download the data if it is not present.
307        voxel_size: The isotropic voxel size in nm of the pyramid level to download. If given, the full raw and
308            label arrays at this level are downloaded. Otherwise, the label level in `DATASETS` is used.
309
310    Returns:
311        The filepath to the cached zarr store.
312    """
313    import zarr
314    from zarr.codecs import BloscCodec
315    from scipy.ndimage import zoom
316
317    if dataset_name not in DATASETS:
318        raise ValueError(f"'{dataset_name}' is not a valid dataset name. Choose from {sorted(DATASETS.keys())}.")
319
320    info = DATASETS[dataset_name]
321    recon, em_name = info["recon"], info["em_name"]
322    label_subpath, label_level = info["label_subpath"], info["label_level"]
323    axis_reverse, full_array = info["axis_reverse"], info["full_array"]
324
325    os.makedirs(str(path), exist_ok=True)
326    zarr_name = dataset_name if voxel_size is None else f"{dataset_name}_{voxel_size:g}nm"
327    zarr_path = os.path.join(str(path), f"{zarr_name}.zarr")
328
329    root = zarr.open_group(zarr_path, mode="a")
330    if "raw" in root and "labels" in root:
331        return zarr_path
332
333    if not download:
334        raise RuntimeError(f"No cached data found at '{zarr_path}'. Set download=True to stream it from S3.")
335
336    print(f"Streaming nucleus data for '{dataset_name}' from the janelia-cosem-datasets S3 bucket ...")
337    label_full_path = f"labels/inference/{label_subpath}"
338    label_levels = _multiscale_levels(f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/{label_full_path}")
339    if voxel_size is None:
340        label_scale = label_levels[label_level]
341        em_level, em_scale = _find_matching_em_level(dataset_name, recon, em_name, label_scale)
342    else:
343        label_level, label_scale = _find_level(label_levels, voxel_size, label_full_path)
344        em_levels = _multiscale_levels(f"{BUCKET_URL}{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}")
345        em_level, em_scale = _find_level(em_levels, voxel_size, f"em/{em_name}")
346        full_array = True
347    resampled = em_level is None
348
349    label_url = (
350        f"s3://janelia-cosem-datasets/{dataset_name}/{dataset_name}.zarr/{recon}/{label_full_path}/{label_level}"
351    )
352    label_arr = _open_remote_zarr(label_url)
353
354    if full_array:
355        label_data = np.asarray(label_arr[:])
356        bounding_box = None
357    else:
358        bounding_box = _auto_bounding_box(
359            dataset_name, recon, label_full_path, label_level, label_scale, CROP_SHAPE_DEFAULT,
360        ) if info.get("bounding_box") is None else info["bounding_box"]
361        label_slices = tuple(slice(*bb) for bb in bounding_box)
362        label_data = label_arr[label_slices]
363
364    if axis_reverse == "label":
365        label_data = np.transpose(label_data, tuple(reversed(range(label_data.ndim))))
366        raw_bbox = tuple(reversed(bounding_box)) if bounding_box is not None else None
367    else:
368        raw_bbox = bounding_box
369
370    if em_level is not None:
371        em_url = f"s3://janelia-cosem-datasets/{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}/{em_level}"
372        em_arr = _open_remote_zarr(em_url)
373        if axis_reverse == "raw":
374            raw_data = np.asarray(em_arr[:])
375            raw_data = np.transpose(raw_data, tuple(reversed(range(raw_data.ndim))))
376            if raw_bbox is not None:
377                raw_data = raw_data[tuple(slice(*bb) for bb in raw_bbox)]
378        elif raw_bbox is not None:
379            raw_data = em_arr[tuple(slice(*bb) for bb in raw_bbox)]
380        else:
381            raw_data = np.asarray(em_arr[:])
382    else:
383        em_level, em_scale = _nearest_em_level(dataset_name, recon, em_name, label_scale)
384        em_url = f"s3://janelia-cosem-datasets/{dataset_name}/{dataset_name}.zarr/{recon}/em/{em_name}/{em_level}"
385        em_arr = _open_remote_zarr(em_url)
386        ratio = [ls / es for ls, es in zip(label_scale, em_scale)]
387        if raw_bbox is not None:
388            em_bbox_native = tuple((int(lo * r), int(hi * r)) for (lo, hi), r in zip(raw_bbox, ratio))
389            raw_native = em_arr[tuple(slice(*bb) for bb in em_bbox_native)]
390        else:
391            raw_native = np.asarray(em_arr[:])
392        zoom_factors = [t / s for t, s in zip(label_data.shape, raw_native.shape)]
393        raw_data = zoom(raw_native, zoom_factors, order=1).astype(raw_native.dtype)
394
395    # Guard against off-by-a-few-voxel shape mismatches between independently stored pyramids.
396    common_shape = tuple(min(r, lb) for r, lb in zip(raw_data.shape, label_data.shape))
397    raw_data = raw_data[tuple(slice(0, s) for s in common_shape)]
398    label_data = label_data[tuple(slice(0, s) for s in common_shape)]
399
400    assert raw_data.shape == label_data.shape, (
401        f"Shape mismatch for '{dataset_name}': raw {raw_data.shape} vs labels {label_data.shape}"
402    )
403
404    def _make_array(name, data, shuffle):
405        arr = root.create_array(
406            name, shape=data.shape, chunks=tuple(min(128, s) for s in data.shape), dtype=data.dtype,
407            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
408        )
409        arr[:] = data
410
411    root.attrs["dataset"] = dataset_name
412    root.attrs["recon"] = recon
413    root.attrs["label_kind"] = info["label_kind"]
414    root.attrs["label_target"] = info["label_target"]
415    root.attrs["label_source"] = label_url
416    root.attrs["em_source"] = em_url
417    root.attrs["resampled"] = resampled
418    root.attrs["voxel_size_nm"] = list(label_scale)
419    root.attrs["license_verified"] = dataset_name not in UNVERIFIED_LICENSE
420
421    _make_array("raw", raw_data, shuffle="shuffle")
422    _make_array("labels", label_data, shuffle="bitshuffle")
423
424    print(f"Cached '{dataset_name}' to '{zarr_path}' (shape {raw_data.shape}).")
425    return zarr_path
426
427
428def get_openorganelle_nucleus_paths(
429    path: Union[os.PathLike, str],
430    dataset_names: Optional[List[str]] = None,
431    download: bool = False,
432    voxel_size: Optional[float] = None,
433) -> List[str]:
434    """Get paths to cached OpenOrganelle nucleus zarr stores.
435
436    Args:
437        path: Filepath to a folder where the cached zarr stores will be saved.
438        dataset_names: The names of the datasets to use. Defaults to all datasets in `DATASETS`.
439        download: Whether to download the data if it is not present.
440        voxel_size: The isotropic voxel size in nm of the pyramid level to download. If given, the full raw and
441            label arrays at this level are downloaded. Otherwise, the label level in `DATASETS` is used.
442
443    Returns:
444        List of filepaths to the cached zarr stores.
445    """
446    if dataset_names is None:
447        dataset_names = list(DATASETS.keys())
448    return [get_openorganelle_nucleus_data(path, name, download, voxel_size) for name in dataset_names]
449
450
451def get_openorganelle_nucleus_dataset(
452    path: Union[os.PathLike, str],
453    patch_shape: Tuple[int, int, int],
454    dataset_names: Optional[List[str]] = None,
455    download: bool = False,
456    offsets: Optional[List[List[int]]] = None,
457    boundaries: bool = False,
458    voxel_size: Optional[float] = None,
459    **kwargs,
460) -> Dataset:
461    """Get the OpenOrganelle nucleus dataset for nucleus segmentation.
462
463    Args:
464        path: Filepath to a folder where the cached zarr stores will be saved.
465        patch_shape: The patch shape (z, y, x) to use for training.
466        dataset_names: The names of the datasets to use. Defaults to all datasets in `DATASETS`.
467        download: Whether to download the data if it is not present.
468        offsets: Offset values for affinity computation used as target.
469        boundaries: Whether to compute boundaries as the target.
470        voxel_size: The isotropic voxel size in nm of the pyramid level to download. If given, the full raw and
471            label arrays at this level are downloaded. Otherwise, the label level in `DATASETS` is used.
472        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
473
474    Returns:
475        The segmentation dataset.
476    """
477    assert len(patch_shape) == 3
478
479    paths = get_openorganelle_nucleus_paths(path, dataset_names, download, voxel_size)
480
481    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
482    kwargs, _ = util.add_instance_label_transform(
483        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
484    )
485
486    return torch_em.default_segmentation_dataset(
487        raw_paths=paths,
488        raw_key="raw",
489        label_paths=paths,
490        label_key="labels",
491        patch_shape=patch_shape,
492        **kwargs,
493    )
494
495
496def get_openorganelle_nucleus_loader(
497    path: Union[os.PathLike, str],
498    patch_shape: Tuple[int, int, int],
499    batch_size: int,
500    dataset_names: Optional[List[str]] = None,
501    download: bool = False,
502    offsets: Optional[List[List[int]]] = None,
503    boundaries: bool = False,
504    voxel_size: Optional[float] = None,
505    **kwargs,
506) -> DataLoader:
507    """Get the DataLoader for nucleus segmentation in the OpenOrganelle nucleus dataset.
508
509    Args:
510        path: Filepath to a folder where the cached zarr stores will be saved.
511        patch_shape: The patch shape (z, y, x) to use for training.
512        batch_size: The batch size for training.
513        dataset_names: The names of the datasets to use. Defaults to all datasets in `DATASETS`.
514        download: Whether to download the data if it is not present.
515        offsets: Offset values for affinity computation used as target.
516        boundaries: Whether to compute boundaries as the target.
517        voxel_size: The isotropic voxel size in nm of the pyramid level to download. If given, the full raw and
518            label arrays at this level are downloaded. Otherwise, the label level in `DATASETS` is used.
519        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
520            or for the PyTorch DataLoader.
521
522    Returns:
523        The DataLoader.
524    """
525    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
526    dataset = get_openorganelle_nucleus_dataset(
527        path, patch_shape, dataset_names=dataset_names, download=download,
528        offsets=offsets, boundaries=boundaries, voxel_size=voxel_size, **ds_kwargs
529    )
530    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)