Merge branch 'main' into remove-fibsem

dacapo/experiments/datasplits/datasets/arrays/zarr_array_config.py

@@ -3,7 +3,7 @@
 from .array_config import ArrayConfig
 
 from funlib.geometry import Coordinate
-from funlib.persistence import open_ds
+from funlib.persistence import open_ds, open_ome_ds
 
 from upath import UPath as Path
 
@@ -56,9 +56,17 @@ class ZarrArrayConfig(ArrayConfig):
     mode: Optional[str] = attr.ib(
         default="a", metadata={"help_text": "The access mode!"}
     )
+    ome_metadata: bool = attr.ib(
+        default=False, metadata={"help_text": "Whether to expect OME metadata"}
+    )
 
     def array(self, mode="r"):
-        return open_ds(f"{self.file_name}/{self.dataset}", mode=mode)
+        if self.ome_metadata:
+            name = self.dataset.split("/")[-1]
+            dataset = self.dataset.replace(f"/{name}", "")
+            return open_ome_ds(self.file_name / dataset, name=name, mode=mode)
+        else:
+            return open_ds(self.file_name / self.dataset, mode=mode)
 
     def verify(self) -> Tuple[bool, str]:
         """

pyproject.toml

@@ -46,7 +46,6 @@ dependencies = [
         "funlib.geometry>=0.2",
         "mwatershed>=0.5.2",
         "cellmap-models",
-        "funlib.persistence>=0.5.3",
         "gunpowder>=1.4",
         "lsds",
         "xarray",
@@ -58,8 +57,8 @@ dependencies = [
         "upath",
         "boto3",
         "matplotlib",
-        "xarray-multiscale",
         "cellmap-schemas",
+        "funlib.persistence @ git+https://github.com/funkelab/funlib.persistence.git@ome-ngff"
         ]
 
 # extras

tests/fixtures/__init__.py

@@ -4,7 +4,7 @@
     unet_architecture,
     unet_3d_architecture,
 )
-from .arrays import dummy_array, zarr_array, cellmap_array
+from .arrays import dummy_array, zarr_array, cellmap_array, multiscale_zarr
 from .datasplits import (
     dummy_datasplit,
     twelve_class_datasplit,

tests/fixtures/arrays.py

@@ -6,7 +6,7 @@
 
 import zarr
 import numpy as np
-
+from numcodecs import Zstd
 import pytest
 
 
@@ -27,7 +27,7 @@ def zarr_array(tmp_path):
         zarr_array_config.dataset, data=np.zeros((100, 50, 25), dtype=np.float32)
     )
     dataset.attrs["offset"] = (12, 12, 12)
-    dataset.attrs["resolution"] = (1, 2, 4)
+    dataset.attrs["voxel_size"] = (1, 2, 4)
     dataset.attrs["axis_names"] = ["z", "y", "x"]
     yield zarr_array_config
 
@@ -45,7 +45,7 @@ def cellmap_array(tmp_path):
         data=np.arange(0, 100, dtype=np.uint8).reshape(10, 5, 2),
     )
     dataset.attrs["offset"] = (12, 12, 12)
-    dataset.attrs["resolution"] = (1, 2, 4)
+    dataset.attrs["voxel_size"] = (1, 2, 4)
     dataset.attrs["axis_names"] = ["z", "y", "x"]
 
     cellmap_array_config = BinarizeArrayConfig(
@@ -59,3 +59,67 @@ def cellmap_array(tmp_path):
     )
 
     yield cellmap_array_config
+
+
+@pytest.fixture()
+def multiscale_zarr(tmp_path):
+    zarr_metadata = {
+        "multiscales": [
+            {
+                "axes": [
+                    {"name": "z", "type": "space", "unit": "nanometer"},
+                    {"name": "y", "type": "space", "unit": "nanometer"},
+                    {"name": "x", "type": "space", "unit": "nanometer"},
+                ],
+                "coordinateTransformations": [],
+                "datasets": [
+                    {
+                        "coordinateTransformations": [
+                            {"scale": [4.2, 7.4, 5.6], "type": "scale"},
+                            {"translation": [6.0, 10.0, 2.0], "type": "translation"},
+                        ],
+                        "path": "s0",
+                    },
+                    {
+                        "coordinateTransformations": [
+                            {"type": "scale", "scale": [1.0, 2.0, 4.0]},
+                            {"type": "translation", "translation": [12.0, 12.0, 12.0]},
+                        ],
+                        "path": "s1",
+                    },
+                ],
+                "name": "multiscale_dataset",
+                "version": "0.4",
+            }
+        ],
+        "omero": {
+            "id": 1,
+            "name": "test_image",
+            "channels": [],
+        },
+    }
+    ome_zarr_array_config = ZarrArrayConfig(
+        name="ome_zarr_array",
+        file_name=tmp_path / "ome_zarr_array.zarr",
+        dataset="multiscale_dataset/s1",
+        ome_metadata=True,
+    )
+
+    store = zarr.DirectoryStore(ome_zarr_array_config.file_name)
+    multiscale_group = zarr.group(
+        store=store, path="multiscale_dataset", overwrite=True
+    )
+
+    for level in [0, 1]:
+        scaling = pow(2, level)
+        multiscale_group.require_dataset(
+            name=f"s{level}",
+            shape=(100 / scaling, 80 / scaling, 60 / scaling),
+            chunks=10,
+            dtype=np.float32,
+            compressor=Zstd(level=6),
+        )
+
+    multiscale_group.attrs.update(zarr_metadata)
+
+    yield ome_zarr_array_config

tests/operations/test_array.py

@@ -0,0 +1,28 @@
+import os
+from ..fixtures import *
+
+import pytest
+from pytest_lazy_fixtures import lf
+
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+
+@pytest.mark.parametrize(
+    "array_config",
+    [
+        lf("zarr_array"),
+        lf("multiscale_zarr"),
+    ],
+)
+def test_array(array_config):
+    array = array_config.array()
+
+    assert array.offset == (12, 12, 12), f"offset is not correct, expected (12, 12, 12), got {array.offset}"
+    assert array.voxel_size == (1, 2, 4), f"resolution is not correct, expected (1, 2, 4), got {array.voxel_size}"
+    assert array.axis_names == ["z", "y", "x"], f"axis names are not correct, expected ['z', 'y', 'x'], got {array.axis_names}"
+
+    # offset = array.attrs["offset"]
+    # resolution = array.attrs["resolution"]
+

tests/operations/test_mini.py

@@ -10,7 +10,7 @@
 from dacapo.experiments import Run
 from dacapo.train import train_run
 from dacapo.validate import validate_run
-
+from dacapo.compute_context import create_compute_context
 import zarr
 
 import pytest
@@ -75,6 +75,9 @@ def test_mini(
         num_iterations=1,
     )
     run = Run(run_config)
+    compute_context = create_compute_context()
+    device = compute_context.device
+    run.model.to(device)
 
     if func == "train":
         train_run(run)