Support reading OME-Zarr 0.5 (#295)

* wip

* wip: 2025-02-27

* Undo reformatting

* Remove unused imports

* Fix miscellaneous failing tests.

* Get rename_well working.

* sort import

* Read MMStack with Zarr v3 (#293)

* use fsspec as a wrapper for ZarrTiffStore

* update deprecated readers and tests

* document the helper function

* add back the binary check after a flag

* use kwarg for file mode in reader

* restore dtype attribute

* restore opening zarr with dask

* stringify path name

* remove redundant dtype

* add rich for zarr v3 print tree

* skip tests that require ome-zarr-py

* fix file list checking
it doesn't matter what the order is

* fix commas

* stringify path names

* detect ome key in reader

* allow multiscales to not have version

* wip: allow reading v0.5 images

* normalize store path

* update docstring

* remove unused import

* remove zarr-python synchronizer
zarr-developers/zarr-python#1596

* allow empty zarr examples

* add example json files

* fixture for example v0.5 HCS dataset

* unify ome container unwrapping

* run PR tests for all target branches

* test with absolute paths
zarr-python 3.0.8 breakage

* avoid big endian test cases

* avoid yanked versions

* fix zarr-python breakage

* revert keyword argument name

* fix version detection after zarr-python breakage

* test with absolute path

* Revert "allow empty zarr examples"

This reverts commit 77146dd.

* Test reading NGFF v0.5 with acquire-zarr

* check binary values

* simplify file mode switch

Co-authored-by: Alan Liddell <alan.c.liddell@gmail.com>

* remove redundant deletion

---------

Co-authored-by: Alan Liddell <alan.c.liddell@gmail.com>
Co-authored-by: Alan Liddell <aliddell@chanzuckerberg.com>

Author: ziw-liu

.github/workflows/pr.yml

@@ -6,7 +6,7 @@ name: lint, style, and tests
 on:
   pull_request:
     branches:
-      - main
+      - "*"
 
 jobs:
   style:

iohub/ngff/models.py

@@ -2,7 +2,7 @@
 
 """
 Data model classes with validation for OME-NGFF metadata.
-Developed against OME-NGFF v0.4 and ome-zarr v0.9
+Developed against OME-NGFF v0.4/0.5.2 and ome-zarr v0.9.
 
 Attributes are 'snake_case' with aliases to match NGFF names in JSON output.
 See https://ngff.openmicroscopy.org/0.4/index.html#naming-style
@@ -219,7 +219,9 @@ class VersionMeta(MetaBase):
     """OME-NGFF spec version. Default is the current version (0.4)."""
 
     # SHOULD
-    version: Literal["0.1", "0.2", "0.3", "0.4"] = "0.4"
+    version: Literal["0.1", "0.2", "0.3", "0.4", "0.5"] | None = Field(
+        default=None, exclude=lambda v: v is None
+    )
 
 
 class MultiScaleMeta(VersionMeta):

iohub/ngff/nodes.py

@@ -16,8 +16,6 @@
 
 import numpy as np
 import zarr.codecs
-import zarr.storage
-from numcodecs import Blosc
 from numpy.typing import ArrayLike, DTypeLike, NDArray
 from pydantic import ValidationError
 from zarr.core.group import normalize_path
@@ -57,27 +55,28 @@ def _pad_shape(shape: tuple[int, ...], target: int = 5):
 
 
 def _open_store(
-    store_path: StrOrBytesPath,
+    store_path: StrOrBytesPath | Path,
     mode: Literal["r", "r+", "a", "w", "w-"],
     version: Literal["0.1", "0.4", "0.5"],
 ):
-    if not os.path.isdir(store_path) and mode in ("r", "r+"):
+    store_path = Path(store_path).resolve()
+    if not store_path.exists() and mode in ("r", "r+"):
         raise FileNotFoundError(
-            f"Dataset directory not found at {store_path}."
+            f"Dataset directory not found at {str(store_path)}."
         )
     if version not in ("0.4", "0.5"):
         _logger.warning(
             "IOHub is only tested against OME-NGFF v0.4 and v0.5. "
             f"Requested version {version} may not work properly."
         )
     try:
-        store = zarr.storage.LocalStore(store_path)
-        root = zarr.open_group(
-            store, mode=mode, zarr_format=(3 if version == "0.5" else 2)
-        )
+        zarr_format = None
+        if mode in ("w", "w-") or (mode == "a" and not store_path.exists()):
+            zarr_format = 3 if version == "0.5" else 2
+        root = zarr.open_group(store_path, mode=mode, zarr_format=zarr_format)
     except Exception as e:
         raise RuntimeError(
-            f"Cannot open Zarr root group at {store_path}"
+            f"Cannot open Zarr root group at {str(store_path)}"
         ) from e
     return root
 
@@ -141,6 +140,11 @@ def zattrs(self):
         Assignments will modify the metadata file."""
         return self._group.attrs
 
+    @property
+    def maybe_wrapped_ome_attrs(self):
+        """Container of OME metadata attributes."""
+        return self.zattrs.get("ome") or self.zattrs
+
     @property
     def version(self):
         """NGFF version"""
@@ -571,12 +575,12 @@ def _set_meta(
             )
             example_image: ImageArray = self[
                 self.metadata.multiscales[0].datasets[0].path
-            ].channels
+            ]
             self._channel_names = list(range(example_image.channels))
 
     def _parse_meta(self):
-        multiscales = self.zattrs.get("multiscales")
-        omero = self.zattrs.get("omero")
+        multiscales = self.maybe_wrapped_ome_attrs.get("multiscales")
+        omero = self.maybe_wrapped_ome_attrs.get("omero")
         if multiscales:
             try:
                 self._set_meta(multiscales=multiscales, omero=omero)
@@ -802,6 +806,7 @@ def _check_shape(self, data_shape: tuple[int]):
             )
 
     def _create_compressor_options(self, chunk_shape: Tuple[int, ...] = None):
+        shuffle = zarr.codecs.BloscShuffle.bitshuffle
         if self._zarr_format == 3:
             return {
                 "codecs": [
@@ -812,17 +817,19 @@ def _create_compressor_options(self, chunk_shape: Tuple[int, ...] = None):
                             zarr.codecs.BloscCodec(
                                 cname="zstd",
                                 clevel=1,
-                                shuffle=Blosc.BITSHUFFLE,
+                                shuffle=shuffle,
                             ),
                         ],
                     )
                 ],
             }
         else:
+            from numcodecs import Blosc
+
             return {
                 "compressor": Blosc(
                     cname="zstd", clevel=1, shuffle=Blosc.BITSHUFFLE
-                ),
+                )
             }
 
     def _create_image_meta(
@@ -1379,7 +1386,7 @@ def __init__(
         )
 
     def _parse_meta(self):
-        if well_group_meta := self.zattrs.get("well"):
+        if well_group_meta := self.maybe_wrapped_ome_attrs.get("well"):
             self.metadata = WellGroupMeta(**well_group_meta)
         else:
             self._warn_invalid_meta()
@@ -1622,7 +1629,7 @@ def __init__(
         )
 
     def _parse_meta(self):
-        if plate_meta := self.zattrs.get("plate"):
+        if plate_meta := self.maybe_wrapped_ome_attrs.get("plate"):
             _logger.debug(f"Loading HCS metadata from file: {plate_meta}")
             self.metadata = PlateMeta(**plate_meta)
         else:
@@ -1930,14 +1937,56 @@ def rename_well(self, old: str, new: str):
         self.dump_meta()
 
 
+def _check_file_mode(
+    store_path: Path,
+    mode: Literal["r", "r+", "a", "w", "w-"],
+    disable_path_checking: bool,
+) -> bool:
+    if mode == "a":
+        mode = "r+" if store_path.exists() else "w-"
+    parse_meta = False
+    if mode in ("r", "r+"):
+        parse_meta = True
+    elif mode == "w-":
+        if store_path.exists():
+            raise FileExistsError(store_path)
+    elif mode == "w":
+        if store_path.exists():
+            if (
+                ".zarr" not in str(store_path.resolve())
+                and not disable_path_checking
+            ):
+                raise ValueError(
+                    "Cannot overwrite a path that does not contain '.zarr', "
+                    "use `disable_path_checking=True` if you are sure that "
+                    f"{store_path} should be overwritten."
+                )
+            _logger.warning(f"Overwriting data at {store_path}")
+    else:
+        raise ValueError(f"Invalid persistence mode '{mode}'.")
+    return parse_meta
+
+
+def _detect_layout(meta_keys: list[str]) -> Literal["fov", "hcs"]:
+    if "plate" in meta_keys:
+        return "hcs"
+    elif "multiscales" in meta_keys:
+        return "fov"
+    else:
+        raise KeyError(
+            "Dataset metadata keys ('plate'/'multiscales') not in "
+            f"the found store metadata keys: {meta_keys}. "
+            "Is this a valid OME-Zarr dataset?"
+        )
+
+
 def open_ome_zarr(
     store_path: StrOrBytesPath | Path,
     layout: Literal["auto", "fov", "hcs", "tiled"] = "auto",
     mode: Literal["r", "r+", "a", "w", "w-"] = "r",
     channel_names: list[str] | None = None,
     axes: list[AxisMeta] | None = None,
     version: Literal["0.1", "0.4", "0.5"] = "0.4",
-    synchronizer: zarr.ThreadSynchronizer | zarr.ProcessSynchronizer = None,
     disable_path_checking: bool = False,
     **kwargs,
 ) -> Plate | Position | TiledPosition:
@@ -2003,42 +2052,17 @@ def open_ome_zarr(
         or :py:class:`iohub.ngff.TiledPosition`)
     """
     store_path = Path(store_path)
-    if mode == "a":
-        mode = ("w-", "r+")[int(store_path.exists())]
-    parse_meta = False
-    if mode in ("r", "r+"):
-        parse_meta = True
-    elif mode == "w-":
-        if store_path.exists():
-            raise FileExistsError(store_path)
-    elif mode == "w":
-        if store_path.exists():
-            if (
-                ".zarr" not in str(store_path.resolve())
-                and not disable_path_checking
-            ):
-                raise ValueError(
-                    "Cannot overwrite a path that does not contain '.zarr', "
-                    "use `disable_path_checking=True` if you are sure that "
-                    f"{store_path} should be overwritten."
-                )
-            _logger.warning(f"Overwriting data at {store_path}")
-    else:
-        raise ValueError(f"Invalid persistence mode '{mode}'.")
+    parse_meta = _check_file_mode(
+        store_path, mode, disable_path_checking=disable_path_checking
+    )
     root = _open_store(store_path, mode, version)
     meta_keys = root.attrs.keys() if parse_meta else []
+    if "ome" in meta_keys:
+        meta_keys = root.attrs["ome"].keys()
+        version = root.attrs["ome"].get("version", version)
     if layout == "auto":
         if parse_meta:
-            if "plate" in meta_keys:
-                layout = "hcs"
-            elif "multiscales" in meta_keys:
-                layout = "fov"
-            else:
-                raise KeyError(
-                    "Dataset metadata keys ('plate'/'multiscales') not in "
-                    f"the found store metadata keys: {meta_keys}. "
-                    "Is this a valid OME-Zarr dataset?"
-                )
+            layout = _detect_layout(meta_keys)
         else:
             raise ValueError(
                 "Store layout must be specified when creating a new dataset."
@@ -2059,5 +2083,6 @@ def open_ome_zarr(
         parse_meta=parse_meta,
         channel_names=channel_names,
         axes=axes,
+        version=version,
         **kwargs,
     )

iohub/reader.py

@@ -24,7 +24,7 @@
 
 
 def _find_ngff_version_in_zarr_group(group: zarr.Group) -> str | None:
-    for key in ["plate", "well"]:
+    for key in ["plate", "well", "ome"]:
         if key in group.attrs:
             if v := group.attrs[key].get("version"):
                 return v
@@ -200,8 +200,8 @@ def print_info(path: StrOrBytesPath, verbose=False):
     path = Path(path).resolve()
     try:
         fmt, extra_info = _infer_format(path)
-        if fmt == "omezarr" and extra_info == "0.4":
-            reader = open_ome_zarr(path, mode="r")
+        if fmt == "omezarr" and extra_info in ("0.4", "0.5"):
+            reader = open_ome_zarr(path, mode="r", version=extra_info)
         else:
             reader = read_images(path, data_type=fmt)
     except (ValueError, RuntimeError):

setup.cfg

@@ -38,7 +38,7 @@ install_requires =
     tifffile>=2024.1.30, <2025.5.21
     natsort>=7.1.1
     ndtiff>=2.2.1
-    zarr>=3.0.0
+    zarr>=3.0.8
     rich
     tqdm
     pillow>=9.4.0
@@ -48,6 +48,7 @@ install_requires =
 
 [options.extras_require]
 dev =
+    acquire-zarr
     black
     flake8
     pytest>=5.0.0