diff --git a/docs/api/plugins/ome_zarr.md b/docs/api/plugins/ome_zarr.md index 09cf10c..6a0a082 100644 --- a/docs/api/plugins/ome_zarr.md +++ b/docs/api/plugins/ome_zarr.md @@ -1,7 +1,26 @@ # OME-ZARR conversion plugin -Write any array or image file to a pyramidal OME-ZARR store. Uses only the -core dependencies for arrays and `.zarr` inputs; reading other file formats -needs the optional `bioio` extra (`pip install "patchworks[bioio]"`). +Write any array or image file to a pyramidal OME-ZARR store, add resolution +levels to an existing store, or store a label image inside an OME-ZARR under +the NGFF `labels/` group. Uses only the core dependencies for arrays and +`.zarr` inputs; reading other file formats needs the optional `bioio` extra +(`pip install "patchworks[bioio]"`). + +Pyramids downsample **X and Y only** — `Z` (and channel/time) are kept at full +resolution, matching anisotropic microscopy stacks. + +## to_ome_zarr ::: patchworks.plugins.ome_zarr.to_ome_zarr + +## add_pyramid + +::: patchworks.plugins.ome_zarr.add_pyramid + +## write_labels + +::: patchworks.plugins.ome_zarr.write_labels + +## register_labels + +::: patchworks.plugins.ome_zarr.register_labels diff --git a/docs/guide/ome_zarr_napari.md b/docs/guide/ome_zarr_napari.md index fb5f337..07247a7 100644 --- a/docs/guide/ome_zarr_napari.md +++ b/docs/guide/ome_zarr_napari.md @@ -3,13 +3,38 @@ Two optional plugins close the loop around `tile_process`: convert any input to a fast, pyramidal OME-ZARR, then inspect the image and its labels in napari. -## Why convert to a pyramidal OME-ZARR? +## Everything in one OME-ZARR (the default) + +When you call `tile_process` on a `.zarr` store **without** `write_to`, the +labels are written **back into that same store** under the NGFF +`labels//` group, as their own multi-scale pyramid: + +```python +from patchworks import tile_process + +# labels land in scan.zarr/labels/labels/ with a pyramid — nothing else needed +tile_process("scan.zarr", fn) +``` + +After this, `scan.zarr` holds both the image and its segmentation, each +pyramidal, in a single NGFF store that napari, Fiji and validators read +natively. Pass `write_to="other.zarr"` to instead write a separate +single-resolution label store, or `output_component="cells"` to name the label +image. + +The label pyramid is built lazily (`da.to_zarr`, streamed chunk by chunk), so +it stays OOM-safe even for terabyte volumes. Control it with `pyramid_levels` +and `pyramid_downscale`. + +## Why a pyramid? A single full-resolution array is slow to browse: every pan or zoom touches the whole plane. A **pyramid** stores progressively downsampled copies, so a viewer -only reads the resolution it needs for the current zoom level. OME-ZARR is the -chunked, cloud-friendly NGFF standard that napari (and Fiji, validators, …) -read natively. +only reads the resolution it needs. Pyramids here downsample **X and Y only** — +`Z` (and channel/time) stay at full resolution, matching anisotropic microscopy +stacks. Downsampling is **strided, nearest-neighbour**: the correct choice for +label images, since interpolating label values would invent objects that never +existed. ## Convert any image to OME-ZARR @@ -24,14 +49,6 @@ from patchworks.plugins.ome_zarr import to_ome_zarr # From a proprietary microscope file (lazy, via bioio): to_ome_zarr("scan.czi", "scan.zarr", n_levels=5) - -# From the labels written by tile_process: -import dask.array as da -to_ome_zarr( - da.from_zarr("labels.zarr", component="labels"), - "labels_pyramid.zarr", - axes="zyx", -) ``` !!! note "Install the readers you need" @@ -39,12 +56,27 @@ to_ome_zarr( catch-all reader (needs a JVM). For speed, add native readers for your formats, e.g. `bioio-ome-tiff`, `bioio-czi`, `bioio-lif`, `bioio-nd2`. -Downsampling uses **strided, nearest-neighbour** subsampling. This is the -correct choice for label images: interpolating label values would invent -objects that never existed. Only the spatial axes (`z`/`y`/`x`) are -downsampled — channel and time axes pass through unchanged. +## Add a pyramid to an existing store -## View the result in napari +Already have a flat (single-resolution) zarr? `add_pyramid` writes the missing +levels in place, lazily: + +```python +from patchworks.plugins.ome_zarr import add_pyramid + +add_pyramid("flat.zarr", base="0", n_levels=5) +``` + +And `write_labels` stores any label array inside an existing OME-ZARR under the +`labels/` group (the same thing `tile_process` does by default): + +```python +from patchworks.plugins.ome_zarr import write_labels + +write_labels("scan.zarr", my_labels, name="nuclei") +``` + +## View image + labels in napari `view_in_napari` opens the image and overlays the labels as a proper *Labels* layer in one call. OME-ZARR pyramids are handed to napari as a lazy multi-scale @@ -53,15 +85,13 @@ list, so even huge stores open instantly and only on-screen data is fetched. ```python from patchworks.plugins.napari import view_in_napari -# image as OME-ZARR, labels as the plain store from tile_process: +# one store holding both image and labels/: +view_in_napari("scan.zarr", labels="scan.zarr/labels/labels") + +# or a separate plain label store written with write_to=: view_in_napari("scan.zarr", labels="labels.zarr") ``` -The label store written by `tile_process` keeps its array under the -`output_component` name (default `"labels"`); `view_in_napari` reads that -component and casts it to `int32` for the Labels layer. Pass -`labels_component=...` if you changed it. - !!! note napari is a GUI-heavy extra and is **not** included in `patchworks[all]`. Install it explicitly: `pip install "patchworks[napari]"`. @@ -70,15 +100,15 @@ component and casts it to `int32` for the Labels layer. Pass ```python from patchworks import tile_process -from patchworks.plugins.ome_zarr import to_ome_zarr from patchworks.plugins.napari import view_in_napari -# 1. segment a large image, streaming labels to disk -tile_process("scan.zarr", fn, write_to="labels.zarr", progress=True) - -# 2. (optional) make a pyramid of the raw image for snappy browsing -to_ome_zarr("scan.zarr", "scan_pyramid.zarr") +# 1. segment — labels are written into scan.zarr with a pyramid, by default +tile_process("scan.zarr", fn, progress=True) -# 3. inspect image + labels together -view_in_napari("scan_pyramid.zarr", labels="labels.zarr") +# 2. inspect image + labels together, straight from the one store +view_in_napari("scan.zarr", labels="scan.zarr/labels/labels") ``` + +Plugging in a different segmentation method is just swapping `fn` — any +callable taking a tile and returning an integer label array works (see the +Cellpose and StarDist examples). diff --git a/pyproject.toml b/pyproject.toml index 617584a..a46f202 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,10 @@ [build-system] -requires = ["hatchling"] +requires = ["hatchling", "hatch-vcs"] build-backend = "hatchling.build" [project] name = "patchworks" -version = "0.2.0" +dynamic = ["version"] description = "Tiled processing of arbitrarily large images with globally consistent labels" readme = "README.md" license = { text = "MIT" } @@ -57,6 +57,9 @@ all = ["patchworks[io,gpu,bioio]", "psutil", "tqdm", "scikit-image"] Homepage = "https://github.com/imcf/patchworks" Issues = "https://github.com/imcf/patchworks/issues" +[tool.hatch.version] +source = "vcs" + [tool.hatch.build.targets.wheel] packages = ["src/patchworks"] diff --git a/src/patchworks/__init__.py b/src/patchworks/__init__.py index a47fada..d8e2b6c 100644 --- a/src/patchworks/__init__.py +++ b/src/patchworks/__init__.py @@ -26,6 +26,9 @@ ... overlap=20, write_to="labels.zarr", progress=True) """ +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as _pkg_version + from ._chunks import auto_overlap, auto_tile_shape, auto_tile_shape_cellpose from ._cluster import make_local_cluster from ._core import tile_process @@ -33,7 +36,10 @@ from ._merge import merge_tile_labels from ._relabel import relabel_sequential_array, relabel_sequential_zarr -__version__ = "0.2.0" +try: + __version__ = _pkg_version("patchworks") +except PackageNotFoundError: # not installed (e.g. running from a checkout) + __version__ = "0+unknown" __all__ = [ "tile_process", "merge_tile_labels", diff --git a/src/patchworks/_core.py b/src/patchworks/_core.py index 1c7d047..56cfa3c 100644 --- a/src/patchworks/_core.py +++ b/src/patchworks/_core.py @@ -52,13 +52,15 @@ def tile_process( tile_shape: Union[ tuple[int, ...], Callable[[tuple, Any], tuple], str, None ] = None, - overlap: int = 0, + overlap: int = 16, channel: int | None = 0, level: int = 0, use_gpu: bool = False, progress: bool = False, write_to: Union[str, Path, None] = None, output_component: str = "labels", + pyramid_levels: int = 5, + pyramid_downscale: int = 2, sequential_labels: bool = False, skip_empty: bool = False, empty_threshold: float | None = None, @@ -100,7 +102,9 @@ def tile_process( Voxels of overlap (halo) added to each tile before *fn* is called, so objects near tile boundaries have enough spatial context to be segmented correctly (Cellpose, StarDist, …). The halo is trimmed off - before merging — the output has the original shape. ``0`` disables it. + before merging — the output has the original shape. Defaults to ``16``; + set it to roughly one object diameter (see ``auto_overlap``) for best + results, or ``0`` to disable. Merging is always **touching-label** based: after the halo is trimmed, labels that touch across a tile boundary are merged into one object. @@ -113,10 +117,23 @@ def tile_process( progress: Show a progress bar during the tile-writing and relabel steps. write_to: - Output zarr store path. When None, an auto-temp store is used and its - path is logged. Pass an explicit path to control the output location. + Explicit output zarr store path. Overrides the default behaviour: the + merged labels are written here as a single-resolution array named + ``output_component`` (no pyramid). When None (default) and *image* is a + ``.zarr`` store, labels are written back into that store under the NGFF + ``labels//`` group with an auto pyramid, so the image + and its segmentation live in one file. When None and *image* is an + array, an auto-temp store is used. output_component: - Array name inside ``write_to``. Default ``"labels"``. + Label name. The array inside ``write_to``, or the NGFF label image name + under ``labels/`` when writing into the input store. Default + ``"labels"``. + pyramid_levels: + Number of resolution levels for the in-store label pyramid (only when + writing into the input ``.zarr``). Default 5. + pyramid_downscale: + Per-level X/Y downsampling factor for that pyramid (Z is kept at full + resolution). Default 2. sequential_labels: Renumber merged labels to a contiguous ``1..N`` range. Default False — labels stay globally unique but gappy (block-encoded), which is fine for @@ -144,9 +161,10 @@ def tile_process( Returns ------- da.Array - Globally relabeled array (int32) backed by ``write_to`` (or an - auto-temp zarr when ``write_to`` is None). Never loads the full volume - into RAM. Call ``.compute()`` yourself only if the result fits in RAM. + Globally relabeled array (int32) backed by the output zarr (the input + store's ``labels//0`` by default, ``write_to`` when given, else an + auto-temp zarr). Never loads the full volume into RAM. Call + ``.compute()`` yourself only if the result fits in RAM. Examples -------- @@ -357,30 +375,57 @@ def _cleanup_stage(): _nw = min(4, os.cpu_count() or 1) + # Default: input is a .zarr store and no explicit write_to → labels go back + # *into* the input store under the NGFF labels// group with an auto + # pyramid, so image + segmentation live in one OME-ZARR. + _into_input = ( + write_to is None + and image_source_path is not None + and image_source_path.endswith(".zarr") + ) + + # The merge always writes its result to a concrete store first. if write_to is not None: - _effective_out = str(write_to) + _merge_out = str(write_to) else: - _effective_out = os.path.join( + _merge_out = os.path.join( tempfile.mkdtemp(prefix="bb_merge_"), "merged.zarr" ) - logger.info( - "write_to not set — merged labels in auto-temp %s", _effective_out - ) zarr_native_merge( stage_path, "staged", - _effective_out, + _merge_out, output_component, n_workers=_nw, show_progress=progress, ) if sequential_labels: logger.info("Relabelling to contiguous ids…") - relabel_sequential_zarr(_effective_out, output_component) + relabel_sequential_zarr(_merge_out, output_component) _cleanup_stage() - # Always return a lazy dask array backed by the output zarr. - # Never load the full volume into RAM here — the merge already materialised - # to disk (auto-temp when write_to=None). Caller can .compute() if needed. - return da.from_zarr(_effective_out, component=output_component) + merged = da.from_zarr(_merge_out, component=output_component) + if not _into_input: + # Lazy dask array backed by the merge store. Never loads the full + # volume into RAM. Caller can .compute() if it fits. + return merged + + # Stream the merged labels into the input store as an NGFF label pyramid, + # then drop the temporary merge store. write_labels uses da.to_zarr, so + # this is chunk-streamed and OOM-safe. + import shutil + + from .plugins.ome_zarr import write_labels + + label_group = write_labels( + image_source_path, + merged, + name=output_component, + n_levels=pyramid_levels, + downscale=pyramid_downscale, + overwrite=True, + ) + shutil.rmtree(os.path.dirname(_merge_out), ignore_errors=True) + logger.info("labels stored in input OME-ZARR under %s", label_group) + return da.from_zarr(label_group, component="0") diff --git a/src/patchworks/plugins/ome_zarr.py b/src/patchworks/plugins/ome_zarr.py index b799efb..a5a40c7 100644 --- a/src/patchworks/plugins/ome_zarr.py +++ b/src/patchworks/plugins/ome_zarr.py @@ -21,6 +21,10 @@ the base support with ``pip install "patchworks[bioio]"`` plus the reader(s) you need. +This module also exposes :func:`add_pyramid`, to add resolution levels to an +existing single-resolution store, and :func:`write_labels`, to store a label +image inside an OME-ZARR under the NGFF ``labels/`` group. + Usage ----- >>> from patchworks.plugins.ome_zarr import to_ome_zarr @@ -28,12 +32,6 @@ >>> # From any microscopy file (lazy, via bioio): >>> to_ome_zarr("scan.czi", "scan.zarr") 'scan.zarr' ->>> ->>> # From the labels produced by tile_process: ->>> import dask.array as da ->>> to_ome_zarr(da.from_zarr("labels.zarr", component="labels"), -... "labels_pyramid.zarr", axes="zyx") -'labels_pyramid.zarr' """ from __future__ import annotations @@ -52,15 +50,17 @@ _NGFF_VERSION = "0.4" _SPATIAL_AXES = frozenset("zyx") -_DEFAULT_ORDER = ( - "tczyx" # axis names assigned to a bare N-D array, from the right -) +# Only X and Y are downsampled when building pyramids; Z is kept at full +# resolution (microscopy stacks are already coarse and anisotropic in Z). +_DOWNSAMPLE_AXES = frozenset("yx") +# axis names assigned to a bare N-D array, taken from the right: +_DEFAULT_ORDER = "tczyx" def _default_axes(ndim: int) -> str: """Assign trailing OME axis names to an unlabelled array. - A 3-D array becomes ``"zyx"``, 4-D ``"czyx"``, 5-D ``"tczyx"``. + A 2-D array becomes ``"yx"``, 3-D ``"zyx"``, 4-D ``"czyx"``. """ if ndim > len(_DEFAULT_ORDER): raise ValueError( @@ -69,6 +69,84 @@ def _default_axes(ndim: int) -> str: return _DEFAULT_ORDER[len(_DEFAULT_ORDER) - ndim :] +def _axis_type(name: str) -> str: + if name in _SPATIAL_AXES: + return "space" + return "time" if name == "t" else "channel" + + +def _axes_meta(axes: str) -> list[dict]: + """NGFF ``axes`` metadata for an axes string.""" + return [{"name": a, "type": _axis_type(a)} for a in axes] + + +def _strides(axes: str, downscale: int) -> tuple[int, ...]: + """Per-axis stride: downsample X/Y only; Z, C and T stay at 1.""" + return tuple(downscale if a in _DOWNSAMPLE_AXES else 1 for a in axes) + + +def _write_pyramid( + arr: da.Array, + axes: str, + group_path: str, + *, + n_levels: int, + downscale: int, + chunks: Union[tuple[int, ...], None], + base_name: str = "0", + write_base: bool = True, +) -> list[dict]: + """Write pyramid levels into *group_path* and return NGFF datasets. + + Level 0 is named *base_name*; deeper levels are ``"1"``, ``"2"``, …. When + *write_base* is False the full-resolution array is assumed to already exist + at ``group_path/base_name`` (used by :func:`add_pyramid`) and only the + downsampled levels are written. + """ + strides = _strides(axes, downscale) + datasets: list[dict] = [] + level = arr.rechunk(chunks) if chunks is not None else arr + for i in range(n_levels): + comp = base_name if i == 0 else str(i) + if i > 0 or write_base: + da.to_zarr(level, group_path, component=comp, overwrite=True) + scale = [float(s**i) for s in strides] + datasets.append( + { + "path": comp, + "coordinateTransformations": [ + {"type": "scale", "scale": scale} + ], + } + ) + logger.info( + "pyramid level %s: shape=%s -> %s", comp, level.shape, group_path + ) + next_shape = tuple(s // st for s, st in zip(level.shape, strides)) + if i + 1 < n_levels and min(next_shape) < 1: + logger.info("stopping pyramid at level %d (next too small)", i) + break + level = level[tuple(slice(None, None, st) for st in strides)] + if chunks is not None: + level = level.rechunk(chunks) + return datasets + + +def _write_multiscales( + group_path: str, axes: str, datasets: list[dict], name: str +) -> None: + """Write NGFF ``multiscales`` metadata onto *group_path*.""" + group = zarr.open_group(group_path, mode="a") + group.attrs["multiscales"] = [ + { + "version": _NGFF_VERSION, + "name": name, + "axes": _axes_meta(axes), + "datasets": datasets, + } + ] + + def _open_bioio(path: str, scene: int) -> tuple[da.Array, str]: """Open *path* with bioio and return a lazy ``(array, axes)`` pair. @@ -83,8 +161,8 @@ def _open_bioio(path: str, scene: int) -> tuple[da.Array, str]: raise ImportError( f"reading {ext} requires bioio. Install it with:\n" " pip install 'patchworks[bioio]'\n" - "plus the matching reader plugin, e.g. bioio-ome-tiff, bioio-czi, " - "bioio-lif, bioio-nd2." + "plus the matching reader plugin, e.g. bioio-ome-tiff, " + "bioio-czi, bioio-lif, bioio-nd2." ) from exc img = BioImage(path) @@ -141,8 +219,8 @@ def to_ome_zarr( *source* may be a dask/NumPy array, a ``.zarr`` store, or any image file readable by bioio (CZI, LIF, ND2, OME-TIFF, …). File inputs are read lazily, and every pyramid level is streamed to disk through dask, so the - full volume never needs to fit in RAM. Only the spatial axes - (``z``/``y``/``x``) are downsampled; channel/time axes are kept intact. + full volume never needs to fit in RAM. Only ``x`` and ``y`` are + downsampled; ``z`` (and channel/time) are kept at full resolution. Parameters ---------- @@ -151,19 +229,17 @@ def to_ome_zarr( out_path : str or Path Destination ``.zarr`` store (a directory). axes : str, optional - One character per array dimension, e.g. ``"zyx"``, ``"cyx"`` or - ``"tczyx"``. ``None`` → inferred from bioio metadata for files, or from - the trailing dimensions for bare arrays. Length must equal the number - of array dimensions. + One character per array dimension, e.g. ``"zyx"`` or ``"cyx"``. + ``None`` → inferred from bioio metadata for files, or from the + trailing dimensions for bare arrays. scene : int, optional Scene index to read from multi-scene files (bioio inputs only). n_levels : int, optional - Maximum number of pyramid levels including full resolution. Fewer - levels are written if a spatial dimension would shrink below 1 px. + Maximum number of pyramid levels including full resolution. downscale : int, optional Per-level downsampling factor along each spatial axis (default 2). chunks : tuple of int, optional - Chunk shape for the written levels. ``None`` keeps the array's chunks. + Chunk shape for the written levels. ``None`` keeps the chunks. overwrite : bool, optional Overwrite an existing store at *out_path*. @@ -172,14 +248,6 @@ def to_ome_zarr( str The path to the written store (``str(out_path)``). - Raises - ------ - ValueError - If ``len(axes)`` does not match the array, ``n_levels < 1`` or - ``downscale < 2``. - ImportError - If a non-zarr file is given but bioio is not installed. - Examples -------- >>> from patchworks.plugins.ome_zarr import to_ome_zarr @@ -197,57 +265,195 @@ def to_ome_zarr( f"axes {axes!r} has {len(axes)} entries but array is {arr.ndim}-D" ) - # Per-axis stride: downsample spatial axes only, leave c/t at stride 1. - strides = tuple(downscale if a in _SPATIAL_AXES else 1 for a in axes) - out = str(out_path) - # Create (or wipe) the group; component arrays are written by dask below. zarr.open_group(out, mode="w" if overwrite else "w-") + datasets = _write_pyramid( + arr, + axes, + out, + n_levels=n_levels, + downscale=downscale, + chunks=chunks, + ) + _write_multiscales(out, axes, datasets, Path(out).stem) + return out - datasets: list[dict] = [] - level = arr.rechunk(chunks) if chunks is not None else arr - for i in range(n_levels): - da.to_zarr(level, out, component=str(i), overwrite=overwrite) - # NGFF scale transform: downscale**i on spatial axes, 1.0 elsewhere. - scale = [float(s**i) for s in strides] - datasets.append( - { - "path": str(i), - "coordinateTransformations": [ - {"type": "scale", "scale": scale} - ], - } + +def add_pyramid( + group_path: Union[str, Path], + *, + base: str = "0", + axes: Union[str, None] = None, + n_levels: int = 5, + downscale: int = 2, + chunks: Union[tuple[int, ...], None] = None, +) -> str: + """Add downsampled pyramid levels to an existing single-resolution zarr. + + Reads the full-resolution array already stored at ``group_path/base``, + writes the missing downsampled levels next to it, and (re)writes the NGFF + ``multiscales`` metadata — turning a flat store into a multi-scale one + in place, lazily. + + Parameters + ---------- + group_path : str or Path + Path to the zarr group holding the base array. + base : str, optional + Component name of the existing full-resolution array (default + ``"0"``). Ignored if the group already has ``multiscales`` metadata, + in which case its level-0 path and axes are reused. + axes : str, optional + Axes string. ``None`` → taken from existing metadata, else inferred + from the array's dimensions. + n_levels, downscale, chunks + As in :func:`to_ome_zarr`. + + Returns + ------- + str + The path to the updated group. + """ + if downscale < 2: + raise ValueError("downscale must be >= 2") + if n_levels < 1: + raise ValueError("n_levels must be >= 1") + + gp = str(group_path) + root = zarr.open_group(gp, mode="r") + multiscales = root.attrs.get("multiscales") + if multiscales: + base = multiscales[0]["datasets"][0]["path"] + if axes is None: + axes = "".join(a["name"] for a in multiscales[0]["axes"]) + + base_arr = da.from_zarr(gp, component=base) + if axes is None: + axes = _default_axes(base_arr.ndim) + if len(axes) != base_arr.ndim: + raise ValueError( + f"axes {axes!r} has {len(axes)} entries but array is " + f"{base_arr.ndim}-D" ) - logger.info("OME-ZARR level %d: shape=%s -> %s", i, level.shape, out) - # Stop once another downsample would erase a spatial dimension. - next_shape = tuple(s // st for s, st in zip(level.shape, strides)) - if i + 1 < n_levels and min(next_shape) < 1: - logger.info( - "stopping pyramid at level %d (next level too small)", i - ) - break - level = level[tuple(slice(None, None, st) for st in strides)] - if chunks is not None: - level = level.rechunk(chunks) + datasets = _write_pyramid( + base_arr, + axes, + gp, + n_levels=n_levels, + downscale=downscale, + chunks=chunks, + base_name=base, + write_base=False, + ) + _write_multiscales(gp, axes, datasets, Path(gp).stem) + return gp + + +def register_labels( + image_store: Union[str, Path], + name: str = "labels", + *, + axes: Union[str, None] = None, + n_levels: int = 5, + downscale: int = 2, + chunks: Union[tuple[int, ...], None] = None, +) -> str: + """Pyramidalise and register an existing ``labels//0`` base level. - root = zarr.open_group(out, mode="a") - root.attrs["multiscales"] = [ - { - "version": _NGFF_VERSION, - "name": Path(out).stem, - "axes": [ - { - "name": a, - "type": "space" - if a in _SPATIAL_AXES - else "time" - if a == "t" - else "channel", - } - for a in axes - ], - "datasets": datasets, - } - ] - return out + Assumes the full-resolution label array already exists at + ``image_store/labels//0`` (e.g. written there directly by + ``tile_process``). Adds the downsampled levels, tags the group with NGFF + ``image-label`` metadata, and lists *name* in ``labels/.zattrs``. + + Returns + ------- + str + Path to the label group (``image_store/labels/``). + """ + store = str(image_store) + group = f"{store}/labels/{name}" + add_pyramid( + group, + base="0", + axes=axes, + n_levels=n_levels, + downscale=downscale, + chunks=chunks, + ) + grp = zarr.open_group(group, mode="a") + grp.attrs["image-label"] = {"version": _NGFF_VERSION} + + labels_grp = zarr.open_group(f"{store}/labels", mode="a") + registered = list(labels_grp.attrs.get("labels", [])) + if name not in registered: + registered.append(name) + labels_grp.attrs["labels"] = registered + return group + + +def write_labels( + image_store: Union[str, Path], + labels: Union[da.Array, np.ndarray], + *, + name: str = "labels", + axes: Union[str, None] = None, + n_levels: int = 5, + downscale: int = 2, + chunks: Union[tuple[int, ...], None] = None, + overwrite: bool = False, +) -> str: + """Store *labels* inside *image_store* under the NGFF ``labels/`` group. + + The labels are written as their own multi-scale pyramid at + ``image_store/labels//`` and registered in + ``image_store/labels/.zattrs``, so the image and its segmentation live in a + single OME-ZARR store (the NGFF *image-label* convention). + + Parameters + ---------- + image_store : str or Path + Existing OME-ZARR store to attach the labels to. + labels : da.Array or np.ndarray + Label array (integer). + name : str, optional + Label image name (default ``"labels"``). + axes, n_levels, downscale, chunks + As in :func:`to_ome_zarr`. + overwrite : bool, optional + Overwrite an existing label image of the same name. + + Returns + ------- + str + Path to the written label group (``image_store/labels/``). + """ + arr = labels if isinstance(labels, da.Array) else da.asarray(labels) + if axes is None: + axes = _default_axes(arr.ndim) + if len(axes) != arr.ndim: + raise ValueError( + f"axes {axes!r} has {len(axes)} entries but array is {arr.ndim}-D" + ) + + store = str(image_store) + # Build the labels/ group hierarchy from the root store so the NGFF + # group markers are persisted at every level (a nested open_group on its + # own does not create the parent `labels` group on a zarr-v3 LocalStore). + root = zarr.open_group(store, mode="a") + parent = root.require_group("labels") + if overwrite and name in parent: + del parent[name] + parent.require_group(name) + + label_group = f"{store}/labels/{name}" + base = arr.rechunk(chunks) if chunks is not None else arr + da.to_zarr(base, label_group, component="0", overwrite=True) + return register_labels( + store, + name, + axes=axes, + n_levels=n_levels, + downscale=downscale, + chunks=chunks, + ) diff --git a/tests/test_core.py b/tests/test_core.py index 6fa016b..c3cb26c 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -101,7 +101,10 @@ def counting_fn(tile): call_count[0] += 1 return _label_fn(tile) - tile_process(arr, counting_fn, skip_empty=True, empty_threshold=0) + # overlap=0 so empty tiles are not fed signal from a neighbour's halo + tile_process( + arr, counting_fn, overlap=0, skip_empty=True, empty_threshold=0 + ) # With staging, fn is called once per non-empty tile assert call_count[0] == 2, f"Expected 2 fn calls, got {call_count[0]}" @@ -154,7 +157,9 @@ def test_merge_transitive_three_tiles(tmp_path): sp = str(tmp_path / "stage.zarr") root = zarr.open_group(sp, mode="w") - a = root.zeros("staged", shape=(3, 4, 4), chunks=(1, 4, 4), dtype=np.int32) + a = root.zeros( + name="staged", shape=(3, 4, 4), chunks=(1, 4, 4), dtype=np.int32 + ) a[0] = np.full((4, 4), 10) # label 10 in tile 0 a[1] = np.full((4, 4), 20) # label 20 in tile 1 (touches 10 and 30) a[2] = np.full((4, 4), 30) # label 30 in tile 2 @@ -176,7 +181,9 @@ def test_merge_isolated_labels_not_merged(tmp_path): sp = str(tmp_path / "stage.zarr") root = zarr.open_group(sp, mode="w") - a = root.zeros("staged", shape=(2, 4, 8), chunks=(1, 4, 4), dtype=np.int32) + a = root.zeros( + name="staged", shape=(2, 4, 8), chunks=(1, 4, 4), dtype=np.int32 + ) # tile (z=0, x-left): label 1 only in left half, no boundary voxel # tile (z=0, x-right): label 2 only in right half # They share the x=4 boundary but fill opposite ends → no touching voxel diff --git a/tests/test_napari.py b/tests/test_napari.py index 31ff755..c1379d7 100644 --- a/tests/test_napari.py +++ b/tests/test_napari.py @@ -17,7 +17,7 @@ def test_resolve_image_multiscale(tmp_path): assert isinstance(out, list) assert len(out) == 3 assert all(isinstance(lvl, da.Array) for lvl in out) - assert out[1].shape == (8, 8, 8) + assert out[1].shape == (16, 8, 8) # Z preserved, only X/Y downsampled def test_resolve_labels_plain_zarr(tmp_path): diff --git a/tests/test_ome_zarr.py b/tests/test_ome_zarr.py index 3bd1401..76dd647 100644 --- a/tests/test_ome_zarr.py +++ b/tests/test_ome_zarr.py @@ -4,8 +4,14 @@ import numpy as np import pytest +import zarr + from patchworks import load_ome_zarr -from patchworks.plugins.ome_zarr import to_ome_zarr +from patchworks.plugins.ome_zarr import ( + add_pyramid, + to_ome_zarr, + write_labels, +) def test_pyramid_roundtrip(tmp_path): @@ -19,12 +25,13 @@ def test_pyramid_roundtrip(tmp_path): l1 = load_ome_zarr(out, channel=None, level=1) l2 = load_ome_zarr(out, channel=None, level=2) + # Z is kept at full resolution; only X/Y are downsampled. assert l0.shape == (8, 8, 8) - assert l1.shape == (4, 4, 4) - assert l2.shape == (2, 2, 2) + assert l1.shape == (8, 4, 4) + assert l2.shape == (8, 2, 2) # Full resolution is byte-identical; downsampling is nearest (label-safe). assert np.array_equal(np.asarray(l0), a) - assert np.array_equal(np.asarray(l1), a[::2, ::2, ::2]) + assert np.array_equal(np.asarray(l1), a[:, ::2, ::2]) def test_non_spatial_axis_not_downsampled(tmp_path): @@ -53,3 +60,36 @@ def test_unreadable_format_without_bioio(tmp_path): to_ome_zarr(str(tmp_path / "scan.czi"), tmp_path / "out.zarr") else: pytest.skip("bioio installed; ImportError path not exercised") + + +def test_add_pyramid_to_flat_store(tmp_path): + """add_pyramid turns a single-array store into a multi-scale one.""" + base = np.arange(8 * 8 * 8, dtype="int32").reshape(8, 8, 8) + store = str(tmp_path / "flat.zarr") + da.to_zarr(da.from_array(base, chunks=(8, 8, 8)), store, component="0") + + add_pyramid(store, base="0", axes="zyx", n_levels=3, downscale=2) + + assert load_ome_zarr(store, channel=None, level=0).shape == (8, 8, 8) + l1 = load_ome_zarr(store, channel=None, level=1) + assert l1.shape == (8, 4, 4) # Z preserved + assert np.array_equal(np.asarray(l1), base[:, ::2, ::2]) + + +def test_write_labels_into_store(tmp_path): + """Labels land under labels// as a registered NGFF pyramid.""" + store = to_ome_zarr( + np.zeros((8, 8, 8), "uint16"), tmp_path / "img.zarr", n_levels=2 + ) + labels = np.ones((8, 8, 8), dtype="int32") + + group = write_labels(store, labels, name="cells", n_levels=2) + + # registered in the parent labels group + labels_grp = zarr.open_group(f"{store}/labels", mode="r") + assert "cells" in labels_grp.attrs["labels"] + # readable as a multi-scale label image with image-label metadata + assert load_ome_zarr(group, channel=None, level=0).shape == (8, 8, 8) + assert load_ome_zarr(group, channel=None, level=1).shape == (8, 4, 4) + lg = zarr.open_group(group, mode="r") + assert lg.attrs["image-label"]["version"]