Create volume annotations (#1340)

* make VolumeLayer non-protected

* add VolumeLayer.edit contextmanager

* format

* lint

* add @pytest.mark.use_proxay
 to test

* format

* rechunk mags with Zarr3Config to have expected volume annotations format

* Enum for VolumeLayerEditMode

* test_edit_volume_annotation, test_save_edited_volume_annotation and test_edited_volume_annotation_upload_download

* format and lint

* edited annotation uploadable

* set zarr3 datatype for wk.Annotation.download

* implement VolumeLayerEditMode.MEMORY

* get voxel_size from nml

* init zip when creating volume layer

* store volume layer in zip of zips and save downloaded annotation volumes to disk

* use volume_layers_root argument of Annotation.download

* update cassettes

* load editing layer into tensorstore, as MemoryFileSystem is not compatible with tensorstore (?)

* typing and format

* avoid tensorstore array memory path collision

* rename nml to skeleton in offline_merger_mode.py

Co-authored-by: Tom Herold <[email protected]>

* replace volume_layers_root with temp dir and incorporate further review

* fix format mistakes

* update cassettes

* extract VolumeLayer to volume_layer.py

* update cassette

* faster VolumeLayer.edit with SequentialExecutor

* improve test_edit_volume_annotation

* document changes

* update learned_segmenter.py to use VolumeLayer.edit

* lint

* VolumeLayer.edit kw-only

---------

Co-authored-by: Tom Herold <[email protected]>

Author: brokkoli71

webknossos/Changelog.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Breaking Changes
 
 ### Added
+- Added context manager `VolumeLayer.edit` for creating and modifying volume annotations. [#1340](https://github.com/scalableminds/webknossos-libs/pull/1340)
 
 ### Changed
 

webknossos/examples/WIP/merge_trees_at_closest_nodes.py

@@ -6,19 +6,19 @@
 
 import webknossos as wk
 
-nml = wk.Skeleton.load("trees-in-groups.nml")
+skeleton = wk.Skeleton.load("trees-in-groups.nml")
 
 # Probably we want to keep groups and normal trees in distinct collections (groups/trees).
 # For many use-cases a common view groups_and_trees would be great, but not here:
-for group in nml.groups.values():  # groups is a dict with the name as keys
+for group in skeleton.groups.values():  # groups is a dict with the name as keys
     min_distance_graph = G = nx.Graph()
     for (tree_idx_a, tree_a), (tree_idx_b, tree_b) in combinations(
         enumerate(group.flattened_trees()), 2
     ):
         pos_a = (
-            tree_a.get_node_positions() * nml.voxel_size
+            tree_a.get_node_positions() * skeleton.voxel_size
         )  # or tree_a.get_node_positions_nm?
-        pos_b = tree_b.get_node_positions() * nml.voxel_size
+        pos_b = tree_b.get_node_positions() * skeleton.voxel_size
         node_idx_a, node_idx_b, distance = wk.geometry.closest_pair(pos_a, pos_b)
         G.add_edge((tree_idx_a, node_idx_a), (tree_idx_b, node_idx_b), weight=distance)
     new_edges = nx.algorithms.tree.mst.minimum_spanning_edges()
@@ -35,7 +35,7 @@
     final_tree.name = group.name
     final_tree.group = None
 
-    del nml.groups[group.name]
+    del skeleton.groups[group.name]
     # or
     group.delete()
     # The latter only works if everything is double-linked.
@@ -44,4 +44,4 @@
     # to do the double-linking. Simply dict-like insertions can't work then:
     # nml["tree-name"] = Tree()
 
-nml.save("merged-trees.nml")
+skeleton.save("merged-trees.nml")

webknossos/examples/WIP/offline_merger_mode.py

@@ -4,21 +4,21 @@
 
 import webknossos as wk
 
-# A merger mode nml with every tree corresponding to a new merged segment is available.
+# A merger mode skeleton with every tree corresponding to a new merged segment is available.
 # All segments in which a node is placed should be merged and saved as a new dataset.
 
-# for local nml:
-nml = wk.open("merger-mode.nml")
+# for local skeleton:
+skeleton = wk.open("merger-mode.skeleton")
 # wk.Skeleton.load or wk.open_skeleton works, too (and is type-safe)
 
 # for online annotation:
-annotation = wk.Annotation.download(
-    "https://webknossos.org/annotations/Explorational/6114d9410100009f0096c640"
-)
-nml = annotation.skeleton
+skeleton = wk.Annotation.download(
+    "https://webknossos.org/annotations/Explorational/6114d9410100009f0096c640",
+    skip_volume_data=True,
+).skeleton
 # should this save anything to disk, or just happen in memory?
 
-dataset = wk.download(nml.dataset_name, organization=nml.dataset_organization)
+dataset = wk.download(skeleton.dataset_name, organization=skeleton.dataset_organization)
 # asks for auth token, persisted into .env or similar config file (maybe use xdg-path?)
 
 # sub-part access via dicts or dict-like classes
@@ -28,7 +28,7 @@
 
 segmentation_data = view.read()
 
-for tree in nml.trees():  # nml.trees() is a flattened iterator of all trees
+for tree in skeleton.trees:  # skeleton.trees() is a flattened iterator of all trees
     segment_ids_in_tree = set(
         segmentation_data[tuple(node.position - view.topleft)] for node in tree.nodes
     )

webknossos/examples/apply_merger_mode.py

@@ -12,8 +12,9 @@ def main() -> None:
     # Opening a merger mode annotation #
     ####################################
 
-    nml = wk.Annotation.download(
-        "https://webknossos.org/annotations/6748612b0100001101c81156"
+    skeleton = wk.Annotation.download(
+        "https://webknossos.org/annotations/6748612b0100001101c81156",
+        skip_volume_data=True,
     ).skeleton
 
     ###############################################
@@ -33,7 +34,7 @@ def main() -> None:
     ##############################
 
     segment_id_mapping = {}
-    for tree in nml.flattened_trees():
+    for tree in skeleton.flattened_trees():
         base = None
         for node in tree.nodes:
             segment_id = in_mag1.read(
@@ -44,7 +45,7 @@ def main() -> None:
             segment_id_mapping[segment_id] = base
 
     print(
-        f"Found {len(list(nml.flattened_trees()))} segment id groups with {len(segment_id_mapping)} nodes"
+        f"Found {len(list(skeleton.flattened_trees()))} segment id groups with {len(segment_id_mapping)} nodes"
     )
     print(segment_id_mapping)
 

webknossos/examples/learned_segmenter.py

@@ -1,6 +1,4 @@
-import os
 from functools import partial
-from tempfile import TemporaryDirectory
 
 import numpy as np
 from skimage import feature
@@ -19,7 +17,6 @@ def main() -> None:
     # Step 1: Read the training data from the annotation and the dataset's color
     # layer (the data will be streamed from WEBKNOSSOS to our local computer)
     training_data_bbox = annotation.user_bounding_boxes[0]  # type: ignore[index]
-    new_dataset_name = f"{annotation.dataset_name.replace(' ', '_')}_segmented"
     with wk.webknossos_context("https://webknossos.org"):
         dataset = annotation.get_remote_annotation_dataset()
 
@@ -58,28 +55,15 @@ def main() -> None:
     assert segmentation.max() < 256
     segmentation = segmentation.astype("uint8")
 
-    # Step 5: Bundle everything as a WEBKNOSSOS layer and upload to wK for viewing and further work
-    with TemporaryDirectory() as tempdir:
-        new_dataset = wk.Dataset(
-            tempdir, voxel_size=dataset.voxel_size, name=new_dataset_name
-        )
-        segmentation_layer = new_dataset.add_layer(
-            "segmentation",
-            wk.SEGMENTATION_CATEGORY,
-            dtype_per_channel=segmentation.dtype,
-            largest_segment_id=int(segmentation.max()),
-        )
+    # Step 5: Upload the segmentation to WEBKNOSSOS
+    print("Uploading segmentation…")
+    volume_layer = annotation.add_volume_layer("segmentation", dtype=segmentation.dtype)
+    with volume_layer.edit() as segmentation_layer:
         segmentation_layer.bounding_box = dataset.layers["color"].bounding_box
         segmentation_layer.add_mag(mag, compress=True).write(segmentation)
         segmentation_layer.downsample(sampling_mode="constant_z")
 
-        remote_ds = new_dataset.upload(
-            layers_to_link=[annotation.get_remote_base_dataset().layers["color"]]
-            if "PYTEST_CURRENT_TEST" not in os.environ
-            else None
-        )
-
-    url = remote_ds.url
+    url = annotation.upload()
     print(f"Successfully uploaded {url}")
 
 

webknossos/examples/skeleton_path_length.py

@@ -6,9 +6,7 @@
 def calculate_path_length(annotation_url: str, auth_token: str) -> None:
     with wk.webknossos_context(token=auth_token):
         # Download a annotation directly from the WEBKNOSSOS server
-        annotation = wk.Annotation.download(
-            annotation_url,
-        )
+        annotation = wk.Annotation.download(annotation_url, skip_volume_data=True)
 
         skeleton = annotation.skeleton
         voxel_size = annotation.voxel_size

webknossos/tests/cassettes/test_annotation/test_edited_volume_annotation_upload_download.yml

@@ -4,8 +4,11 @@
 
 import numpy as np
 import pytest
+from cluster_tools import get_executor
 
 import webknossos as wk
+from webknossos import Annotation, SegmentationLayer
+from webknossos.annotation.volume_layer import VolumeLayerEditMode
 from webknossos.dataset import DataFormat
 from webknossos.geometry import BoundingBox, Vec3Int
 
@@ -40,7 +43,10 @@ def test_annotation_from_wkw_zip_file() -> None:
     assert len(list(copied_annotation.get_volume_layer_names())) == 1
     assert len(list(copied_annotation.skeleton.flattened_trees())) == 1
 
-    copied_annotation.add_volume_layer(name="new_volume_layer")
+    copied_annotation.add_volume_layer(
+        name="new_volume_layer",
+        dtype=np.uint32,
+    )
     assert len(list(copied_annotation.get_volume_layer_names())) == 2
     copied_annotation.delete_volume_layer(volume_layer_name="new_volume_layer")
     assert len(list(copied_annotation.get_volume_layer_names())) == 1
@@ -368,3 +374,158 @@ def test_tree_metadata(tmp_path: Path) -> None:
         list(tmp_annotation.skeleton.flattened_trees())[0].metadata["test_tree"]
         == "test"
     )
+
+
+@pytest.mark.parametrize(
+    "edit_mode", [VolumeLayerEditMode.MEMORY, VolumeLayerEditMode.TEMPORARY_DIRECTORY]
+)
+@pytest.mark.parametrize("executor", ["sequential", "multiprocessing"])
+def test_edit_volume_annotation(edit_mode: VolumeLayerEditMode, executor: str) -> None:
+    dtype = np.uint32
+    data = np.ones((1, 10, 10, 10), dtype=dtype)
+    ann = wk.Annotation(
+        name="my_annotation",
+        dataset_name="sample_dataset",
+        voxel_size=(11.2, 11.2, 25.0),
+    )
+
+    volume_layer = ann.add_volume_layer(
+        name="segmentation",
+        dtype=dtype,
+    )
+    if edit_mode == VolumeLayerEditMode.MEMORY and executor == "multiprocessing":
+        with pytest.raises(ValueError, match="SequentialExecutor"):
+            with volume_layer.edit(
+                edit_mode=edit_mode, executor=get_executor(executor)
+            ) as seg_layer:
+                pass
+    else:
+        with volume_layer.edit(
+            edit_mode=edit_mode, executor=get_executor(executor)
+        ) as seg_layer:
+            assert isinstance(seg_layer, SegmentationLayer)
+            mag = seg_layer.add_mag(1)
+            mag.write(data, absolute_offset=(0, 0, 0), allow_resize=True)
+        with volume_layer.edit(edit_mode=edit_mode) as seg_layer:
+            assert len(seg_layer.mags) == 1
+            mag = seg_layer.get_mag(1)
+            read_data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+            assert np.array_equal(data, read_data)
+
+
+def test_edited_volume_annotation_format() -> None:
+    import zipfile
+
+    import tensorstore
+
+    path = TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
+    ann = Annotation.load(path)
+    data = np.ones(shape=(10, 10, 10))
+
+    volume_layer = ann.add_volume_layer(
+        name="segmentation",
+        dtype=np.uint32,
+    )
+    with volume_layer.edit() as seg_layer:
+        mag_view = seg_layer.add_mag(1)
+        mag_view.write(data, allow_resize=True)
+
+    save_path = TESTOUTPUT_DIR / "saved_annotation.zip"
+    ann.save(save_path)
+    unpack_dir = TESTOUTPUT_DIR / "unpacked_annotation"
+    with zipfile.ZipFile(save_path, "r") as zip_ref:
+        zip_ref.extractall(unpack_dir)
+
+    # test for the format assumptions as mentioned in https://github.com/scalableminds/webknossos/issues/8604
+    ts = tensorstore.open(
+        {
+            "driver": "zarr3",
+            "kvstore": {
+                "driver": "zip",
+                "path": "volumeAnnotationData/1/",
+                "base": {
+                    "driver": "file",
+                    "path": str(unpack_dir / "data_1_segmentation.zip"),
+                },
+            },
+        },
+        create=False,
+        open=True,
+    ).result()
+    metadata = ts.spec().to_json()["metadata"]
+
+    assert metadata["chunk_key_encoding"] == {
+        "configuration": {"separator": "."},
+        "name": "default",
+    }
+    assert ["transpose", "bytes", "blosc"] == [
+        codec["name"] for codec in metadata["codecs"]
+    ]
+    data_read = ts.read().result()[0, :10, :10, :10]
+    assert np.array_equal(data, data_read)
+
+
+@pytest.mark.parametrize(
+    "edit_mode", [VolumeLayerEditMode.MEMORY, VolumeLayerEditMode.TEMPORARY_DIRECTORY]
+)
+def test_edited_volume_annotation_save_load(edit_mode: VolumeLayerEditMode) -> None:
+    data = np.ones((1, 10, 10, 10))
+
+    ann = wk.Annotation(
+        name="my_annotation",
+        dataset_name="sample_dataset",
+        voxel_size=(11.2, 11.2, 25.0),
+    )
+
+    volume_layer = ann.add_volume_layer(name="segmentation", dtype=np.uint32)
+    with volume_layer.edit(edit_mode=edit_mode) as seg_layer:
+        mag_view = seg_layer.add_mag(1)
+        mag_view.write(data, allow_resize=True)
+
+    save_path = TESTOUTPUT_DIR / "annotation_saved.zip"
+    ann.save(save_path)
+    ann_loaded = Annotation.load(save_path)
+
+    volume_layer_downloaded = ann_loaded.get_volume_layer("segmentation")
+
+    with volume_layer_downloaded.edit(edit_mode=edit_mode) as seg_layer:
+        assert len(seg_layer.mags) == 1
+        mag = seg_layer.get_mag(1)
+        read_data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+        assert np.array_equal(data, read_data)
+
+
+@pytest.mark.use_proxay
+def test_edited_volume_annotation_upload_download() -> None:
+    data = np.ones((1, 10, 10, 10))
+
+    ann = Annotation.load(
+        TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
+    )
+    ann.organization_id = "Organization_X"
+
+    volume_layer = ann.add_volume_layer(
+        name="segmentation",
+        dtype=np.uint32,
+    )
+    with volume_layer.edit() as seg_layer:
+        mag_view = seg_layer.add_mag(1)
+        mag_view.write(data, allow_resize=True)
+
+    url = ann.upload()
+    ann_downloaded = Annotation.download(
+        url,
+    )
+
+    assert {layer.name for layer in ann_downloaded._volume_layers} == {
+        "Volume",
+        "segmentation",
+    }
+
+    volume_layer_downloaded = ann_downloaded.get_volume_layer("segmentation")
+
+    with volume_layer_downloaded.edit() as seg_layer:
+        assert len(seg_layer.mags) == 1
+        mag = seg_layer.get_mag(1)
+        read_data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+        assert np.array_equal(data, read_data)

webknossos/tests/test_annotation.py

@@ -660,14 +660,15 @@ def test_merge_fallback_no_fallback_layer(
         )
 
         annotation._volume_layers = [
-            webknossos.annotation._VolumeLayer(  # type: ignore
+            webknossos.annotation.VolumeLayer(  # type: ignore
                 id=0,
                 name=tmp_layer.name,
                 fallback_layer_name=fallback_mag.layer.name,
                 zip=volume_layer_zip,
                 segments={},
                 data_format=DataFormat.WKW,
                 largest_segment_id=largest_segment_id,
+                voxel_size=tmp_dataset.voxel_size,
             ),
         ]