Merge pull request #82 from will-moore/labels_chunks_nested

Use nested chunk store for labels

setup.py

@@ -35,7 +35,7 @@ def get_long_description() -> str:
     author="The Open Microscopy Team",
     author_email="",
     python_requires=">=3",
-    install_requires=["omero-py>=5.6.0", "ome-zarr"],
+    install_requires=["omero-py>=5.6.0", "ome-zarr>=0.2.0"],
     long_description=long_description,
     keywords=["OMERO.CLI", "plugin"],
     url="https://github.com/ome/omero-cli-zarr/",

src/omero_zarr/masks.py

@@ -16,9 +16,9 @@
 from omero.model import MaskI, PolygonI
 from omero.rtypes import unwrap
 from skimage.draw import polygon as sk_polygon
-from zarr.convenience import open as zarr_open
+from zarr.hierarchy import open_group
 
-from .util import print_status
+from .util import open_store, print_status
 
 # Mapping of dimension names to axes in the Zarr
 DIMENSION_ORDER: Dict[str, int] = {
@@ -278,7 +278,8 @@ def save(self, masks: List[omero.model.Shape], name: str) -> None:
         assert input_pyramid.load(Multiscales), "No multiscales metadata found"
         input_pyramid_levels = len(input_pyramid.data)
 
-        root = zarr_open(filename)
+        store = open_store(filename)
+        root = open_group(store)
 
         if current_path in root.group_keys():
             out_labels = getattr(root, current_path)
@@ -302,17 +303,30 @@ def save(self, masks: List[omero.model.Shape], name: str) -> None:
             ignored_dimensions,
             check_overlaps=True,
         )
+        # For v0.3 ngff we want to reduce the number of dimensions to
+        # match the dims of the Image.
+        dims_to_squeeze = []
+        axes = []
+        for dim, size in enumerate(self.image_shape):
+            if size == 1:
+                dims_to_squeeze.append(dim)
+            else:
+                axes.append("tczyx"[dim])
+        labels = np.squeeze(labels, axis=tuple(dims_to_squeeze))
+
         scaler = Scaler(max_layer=input_pyramid_levels)
         label_pyramid = scaler.nearest(labels)
         pyramid_grp = out_labels.require_group(name)
 
-        write_multiscale(label_pyramid, pyramid_grp)  # TODO: dtype, chunks, overwite
+        write_multiscale(
+            label_pyramid, pyramid_grp, axes=axes
+        )  # TODO: dtype, chunks, overwite
 
         # Specify and store metadata
         image_label_colors: List[JSONDict] = []
         label_properties: List[JSONDict] = []
         image_label = {
-            "version": "0.2",
+            "version": "0.3",
             "colors": image_label_colors,
             "properties": label_properties,
             "source": {"image": source_image_link},

src/omero_zarr/raw_pixels.py

@@ -10,24 +10,10 @@
 from omero.rtypes import unwrap
 from skimage.transform import resize
 from zarr.hierarchy import Array, Group, open_group
-from zarr.storage import FSStore
 
 from . import __version__
 from . import ngff_version as VERSION
-from .util import print_status
-
-
-def _open_store(name: str) -> FSStore:
-    """
-    Create an FSStore instance that supports nested storage of chunks.
-    """
-    return FSStore(
-        name,
-        auto_mkdir=True,
-        key_separator="/",
-        normalize_keys=False,
-        mode="w",
-    )
+from .util import open_store, print_status
 
 
 def image_to_zarr(image: omero.gateway.ImageWrapper, args: argparse.Namespace) -> None:
@@ -36,7 +22,7 @@ def image_to_zarr(image: omero.gateway.ImageWrapper, args: argparse.Namespace) -
 
     name = os.path.join(target_dir, "%s.zarr" % image.id)
     print(f"Exporting to {name} ({VERSION})")
-    store = _open_store(name)
+    store = open_store(name)
     root = open_group(store)
     n_levels, axes = add_image(image, root, cache_dir=cache_dir)
     add_multiscales_metadata(root, axes, n_levels)
@@ -226,7 +212,7 @@ def plate_to_zarr(plate: omero.gateway._PlateWrapper, args: argparse.Namespace)
     target_dir = args.output
     cache_dir = target_dir if args.cache_numpy else None
     name = os.path.join(target_dir, "%s.zarr" % plate.id)
-    store = _open_store(name)
+    store = open_store(name)
     print(f"Exporting to {name} ({VERSION})")
     root = open_group(store)
 

src/omero_zarr/util.py

@@ -1,5 +1,7 @@
 import time
 
+from zarr.storage import FSStore
+
 
 def print_status(t0: int, t: int, count: int, total: int) -> None:
     """Prints percent done and ETA.
@@ -18,3 +20,16 @@ def print_status(t0: int, t: int, count: int, total: int) -> None:
         eta = "NA"
     status = f"{percent_done:.2f}% done, ETA: {eta}"
     print(status, end="\r", flush=True)
+
+
+def open_store(name: str) -> FSStore:
+    """
+    Create an FSStore instance that supports nested storage of chunks.
+    """
+    return FSStore(
+        name,
+        auto_mkdir=True,
+        key_separator="/",
+        normalize_keys=False,
+        mode="w",
+    )