Provide direct access to an underlying Zarr array (#792)

* Provide direct access to an underlying Zarr array

* fixes + tests + changelog

* tests

* Update webknossos/webknossos/dataset/mag_view.py

Co-authored-by: Jonathan Striebel <[email protected]>

* adds zarr+dask example

* simpler dask sample

* doc

* use to_slices

* test

* exclude from vcr

Co-authored-by: Jonathan Striebel <[email protected]>

Author: normanrz

docs/mkdocs.yml

@@ -76,6 +76,7 @@ nav:
         - webknossos-py/examples/upload_image_data.md
         - webknossos-py/examples/download_image_data.md
         - webknossos-py/examples/remote_datasets.md
+        - webknossos-py/examples/zarr_and_dask.md
       - Annotation Examples:
         - webknossos-py/examples/apply_merger_mode.md
         - webknossos-py/examples/learned_segmenter.md

docs/src/webknossos-py/examples/zarr_and_dask.md

@@ -0,0 +1,9 @@
+# Zarr and Dask interoperability
+
+This example shows how to access the underlying [Zarr](https://zarr.dev) array of a [remote datasets](../../api/webknossos/dataset/dataset.md#RemoteDataset). Using the Zarr array allows to use other libraries, such as [Dask](https://www.dask.org/) for parallel processing.
+
+```python
+--8<--
+webknossos/examples/zarr_and_dask.py
+--8<--
+```

webknossos/Changelog.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 ### Breaking Changes
 
 ### Added
+- Added direct access to an underlying Zarr array with the `MagView.get_zarr_array()` method. [#792](https://github.com/scalableminds/webknossos-libs/pull/792)
 
 ### Changed
 

webknossos/examples/zarr_and_dask.py

@@ -0,0 +1,25 @@
+import dask.array as da
+
+import webknossos as wk
+
+
+def main() -> None:
+    # Remote datasets are read-only, but can be used similar to normal datasets:
+    l4_sample_dataset = wk.Dataset.open_remote(
+        "https://webknossos.org/datasets/scalable_minds/l4_sample"
+    )
+
+    layer = l4_sample_dataset.get_layer("color")
+    mag = layer.get_finest_mag()
+
+    zarr_array = mag.get_zarr_array()
+    dask_array = da.from_array(zarr_array, chunks=(1, 256, 256, 256))[
+        (0,) + layer.bounding_box.to_slices()
+    ]
+
+    mean_value = dask_array.mean().compute()
+    print("Mean:", mean_value)
+
+
+if __name__ == "__main__":
+    main()

webknossos/poetry.lock

@@ -70,6 +70,7 @@ all = ["pims","tifffile","imagecodecs","JPype1",]
 [tool.poetry.dev-dependencies]
 # autoflake
 black = "21.7b0"
+dask = "2022.02.0"
 hypothesis = "^6.35.0"
 icecream = "^2.1.1"
 inducoapi = "2.0.0"

webknossos/pyproject.toml

@@ -421,6 +421,26 @@ def test_view_write(data_format: DataFormat, output_path: Path) -> None:
     assert np.array_equal(data, write_data)
 
 
+@pytest.mark.parametrize("output_path", [TESTOUTPUT_DIR, REMOTE_TESTOUTPUT_DIR])
+def test_direct_zarr_access(output_path: Path) -> None:
+    ds_path = copy_simple_dataset(DataFormat.Zarr, output_path)
+    mag = Dataset.open(ds_path).get_layer("color").get_mag("1")
+
+    np.random.seed(1234)
+
+    # write: zarr, read: wk
+    write_data = (np.random.rand(3, 10, 10, 10) * 255).astype(np.uint8)
+    mag.get_zarr_array()[:, 0:10, 0:10, 0:10] = write_data
+    data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    assert np.array_equal(data, write_data)
+
+    # write: wk, read: zarr
+    write_data = (np.random.rand(3, 10, 10, 10) * 255).astype(np.uint8)
+    mag.write(write_data, absolute_offset=(0, 0, 0))
+    data = mag.get_zarr_array()[:, 0:10, 0:10, 0:10]
+    assert np.array_equal(data, write_data)
+
+
 @pytest.mark.parametrize("data_format,output_path", DATA_FORMATS_AND_OUTPUT_PATHS)
 def test_view_write_out_of_bounds(data_format: DataFormat, output_path: Path) -> None:
     ds_path = copy_simple_dataset(

webknossos/tests/test_dataset.py

@@ -222,3 +222,13 @@ def test_remote_datasets() -> None:
     assert ds.url == "http://localhost:9000/datasets/Organization_X/e2006_knossos"
     ds.tags = ["test"]
     assert ds in wk.Dataset.get_remote_datasets(tags=["test"]).values()
+
+
+@pytest.mark.block_network(allowed_hosts=[".*"])
+@pytest.mark.vcr(ignore_hosts=["webknossos.org", "data-humerus.webknossos.org"])
+def test_zarr_and_dask() -> None:
+    import examples.zarr_and_dask as example
+
+    (mean_value,) = exec_main_and_get_vars(example, "mean_value")
+
+    assert 124 < mean_value < 125

webknossos/tests/test_examples.py

@@ -7,11 +7,12 @@
 from uuid import uuid4
 
 import numpy as np
+import zarr
 from upath import UPath
 
 from ..geometry import BoundingBox, Mag, Vec3Int, Vec3IntLike
 from ..utils import get_executor_for_args, is_fs_path, rmtree, wait_and_ensure_success
-from ._array import ArrayInfo, BaseArray
+from ._array import ArrayInfo, BaseArray, ZarrArray
 from .properties import MagViewProperties
 
 if TYPE_CHECKING:
@@ -121,6 +122,15 @@ def path(self) -> Path:
     def name(self) -> str:
         return self._mag.to_layer_name()
 
+    def get_zarr_array(self) -> zarr.Array:
+        """
+        Directly access the underlying Zarr array. Only available for Zarr-based datasets.
+        """
+        array_wrapper = self._array
+        if not isinstance(array_wrapper, ZarrArray):
+            raise ValueError("Cannot get the zarr array for wkw datasets.")
+        return array_wrapper._zarray
+
     def write(
         self,
         data: np.ndarray,