Use ome-zarr-models for writing metadata

docs/tutorial/tutorial.py

@@ -16,6 +16,7 @@
 import skimage.data
 import tifffile
 from loguru import logger
+from pydantic_zarr.v2 import ArraySpec
 
 import stack_to_chunk
 
@@ -69,11 +70,13 @@
 # Once we've created it, the ``levels`` property shows that no levels have been added
 # to the group yet.
 
+
 group = stack_to_chunk.MultiScaleGroup(
     temp_dir_path / "chunked.ome.zarr",
     name="my_zarr_group",
     spatial_unit="centimeter",
     voxel_size=(3, 4, 5),
+    array_spec=ArraySpec.from_array(images, chunks=(16, 16, 16)),
 )
 print(group.levels)
 
@@ -95,7 +98,6 @@
 # %%
 # And finally, lets create our first data copy:
 
-group.create_initial_dataset(images, chunk_size=16, compressor="default")
 group.add_full_res_data(images, n_processes=1)
 
 # %%

pyproject.toml

@@ -16,6 +16,8 @@ dependencies = [
     "joblib==1.4.2",
     "loguru==0.7.3",
     "numpy==2.2.1",
+    "ome-zarr-models@git+https://github.com/dstansby/ome-zarr-models-py@993bd6a2b1de3a3aa274ab3118ffce82ece5f46d",
+    "pydantic-zarr==0.7.0",
     "scikit-image==0.25.0",
     "zarr==2.18.4",
 ]
@@ -24,6 +26,7 @@ dynamic = ["version"]
 keywords = []
 name = "stack_to_chunk"
 optional-dependencies = {dev = [
+    "mypy",
     "pre-commit",
     "tox>=4",
 ], docs = [
@@ -59,9 +62,4 @@ overrides."project.classifiers".inline_arrays = false
 overrides."tool.coverage.paths.source".inline_arrays = false
 
 [tool.uv]
-dev-dependencies = [
-    "ome-zarr-models@git+https://github.com/BioImageTools/ome-zarr-models-py@0068c1e4fcca35a942260f82a299e05ebb180019",
-    "pre-commit>=3.8.0",
-    "pytest-cov>=5.0.0",
-    "pytest>=8.3.2",
-]
+dev-dependencies = ["pre-commit>=3.8.0", "pytest-cov>=5.0.0", "pytest>=8.3.2"]

src/stack_to_chunk/main.py

@@ -4,15 +4,16 @@
 
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Literal
 
 import numpy as np
 import zarr
 from dask.array.core import Array
 from joblib import Parallel
 from loguru import logger
 from numcodecs import blosc
-from numcodecs.abc import Codec
+from ome_zarr_models.v04 import Image
+from ome_zarr_models.v04.axes import Axis
+from pydantic_zarr.v2 import ArraySpec
 
 from stack_to_chunk._array_helpers import _copy_slab, _downsample_block
 from stack_to_chunk.ome_ngff import SPATIAL_UNIT, DatasetDict
@@ -32,7 +33,7 @@ def memory_per_process(input_data: Array, *, chunk_size: int) -> int:
 
 class MultiScaleGroup:
     """
-    A class for creating and interacting with a OME-zarr multi-scale group.
+    A class for creating and interacting with a OME-Zarr multi-scale group.
 
     Parameters
     ----------
@@ -44,6 +45,9 @@ class MultiScaleGroup:
         Size of a single voxel, in units of spatial_units.
     spatial_units :
         Units of the voxel size.
+    array_spec :
+        Specification for initial dataset array. If opening an existing group
+        does not need to be provided.
 
     """
 
@@ -54,42 +58,54 @@ def __init__(
         name: str,
         voxel_size: tuple[float, float, float],
         spatial_unit: SPATIAL_UNIT,
+        array_spec: ArraySpec | None = None,
     ) -> None:
+        self._store = zarr.DirectoryStore(path)
         self._path = path
         self._name = name
         self._spatial_unit = spatial_unit
         self._voxel_size = voxel_size
 
         if isinstance(path, Path) and not path.exists():
-            self._create_zarr_group()
+            if array_spec is None:
+                msg = "Group does not already exist, array_spec must be provided"
+                raise ValueError(msg)
+            self._create_zarr_group(array_spec)
 
-        self._group = zarr.open_group(store=self._path, mode="r+")
+        self._group = zarr.open_group(store=self._store, mode="r+")
 
-    def _create_zarr_group(self) -> None:
+    def _create_zarr_group(self, array_spec: ArraySpec) -> None:
         """
         Create the zarr group.
 
         Saves a reference to the group on the ._group attribute.
         """
-        self._group = zarr.open_group(store=self._path, mode="w")
-        multiscales: dict[str, Any] = {}
-        multiscales["version"] = "0.4"
-        multiscales["name"] = self._name
-        multiscales["axes"] = [
-            {"name": "x", "type": "space", "unit": self._spatial_unit},
-            {"name": "y", "type": "space", "unit": self._spatial_unit},
-            {"name": "z", "type": "space", "unit": self._spatial_unit},
-        ]
-        multiscales["type"] = "local mean"
-        multiscales["metadata"] = {
-            "description": "Downscaled using local mean in 2x2x2 blocks.",
-            "method": "skimage.measure.block_reduce",
-            "version": "0.24.0",
-            "kwargs": {"block_size": 2, "func": "np.mean"},
-        }
-
-        multiscales["datasets"] = []
-        self._group.attrs["multiscales"] = [multiscales]
+        self._image: Image = Image.new(
+            array_specs=[array_spec],
+            paths=["0"],
+            axes=[
+                Axis(name="x", type="space", unit=self._spatial_unit),
+                Axis(name="y", type="space", unit=self._spatial_unit),
+                Axis(name="z", type="space", unit=self._spatial_unit),
+            ],
+            name=self._name,
+            multiscale_type="local mean",
+            metadata={
+                "description": "Downscaled using local mean in 2x2x2 blocks.",
+                "method": "skimage.measure.block_reduce",
+                "version": "0.24.0",
+                "kwargs": {"block_size": 2, "func": "np.mean"},
+            },
+            scales=[self._voxel_size],
+            translations=[
+                (
+                    self._voxel_size[0] / 2,
+                    self._voxel_size[1] / 2,
+                    self._voxel_size[2] / 2,
+                )
+            ],
+        )
+        self._image.to_zarr(store=self._store, path="/")
 
     @property
     def levels(self) -> list[int]:
@@ -111,41 +127,6 @@ def __getitem__(self, level: int) -> zarr.Array:
 
         return self._group[str(level)]
 
-    def create_initial_dataset(
-        self, data: Array, *, chunk_size: int, compressor: Literal["default"] | Codec
-    ) -> None:
-        """
-        Set up the inital full-resolution dataset.
-
-        Parameters
-        ----------
-        data :
-            Full input data. Must be 3D, and have a chunksize of ``(nx, ny, 1)``, where
-            ``(nx, ny)`` is the shape of the input 2D slices.
-        chunk_size :
-            Size of chunks in output zarr dataset.
-        compressor :
-            Compressor to use when writing data to zarr dataset.
-
-        Raises
-        ------
-        RuntimeError :
-            If full resolution dataset has already been created.
-
-        """
-        if "0" in self._group:
-            msg = "Full resolution dataset already set up."
-            raise RuntimeError(msg)
-        self._group.create_dataset(
-            name="0",
-            shape=data.shape,
-            chunks=(chunk_size, chunk_size, chunk_size),
-            dtype=data.dtype,
-            compressor=compressor,
-            dimension_separator="/",
-        )
-        self._add_level_metadata(0)
-
     def add_full_res_data(
         self,
         data: Array,