Repository: scalableminds/webknossos-libs

Issue:
Move functional parts of merge volume annotation CLI to Dataset/Annotation classes
## Detailed Description
- Affected library: webknossos
- Functionality is currently implemented in the CLI merge_fallback.py file. It should be moved to a well fitting place within the libs.



Discussion:

Resolving patch:
diff --git a/webknossos/webknossos/annotation/annotation.py b/webknossos/webknossos/annotation/annotation.py
--- a/webknossos/webknossos/annotation/annotation.py
+++ b/webknossos/webknossos/annotation/annotation.py
@@ -19,6 +19,7 @@
 """
 
 import json
+import logging
 import re
 import warnings
 from contextlib import contextmanager, nullcontext
@@ -45,6 +46,7 @@
 from zlib import Z_BEST_SPEED
 
 import attr
+from cluster_tools.executor_protocol import Executor
 from upath import UPath
 from zipp import Path as ZipPath
 
@@ -62,9 +64,11 @@
 from ..dataset.properties import DatasetProperties, dataset_converter
 from ..geometry import NDBoundingBox, Vec3Int
 from ..skeleton import Skeleton
-from ..utils import time_since_epoch_in_ms, warn_deprecated
+from ..utils import get_executor_for_args, time_since_epoch_in_ms, warn_deprecated
 from ._nml_conversion import annotation_to_nml, nml_to_skeleton
 
+logger = logging.getLogger(__name__)
+
 Vector3 = Tuple[float, float, float]
 Vector4 = Tuple[float, float, float, float]
 
@@ -440,9 +444,11 @@ def _parse_volumes(
                     id=volume.id,
                     name="Volume" if volume.name is None else volume.name,
                     fallback_layer_name=volume.fallback_layer,
-                    data_format=DataFormat(volume.format)
-                    if volume.format is not None
-                    else DataFormat.WKW,
+                    data_format=(
+                        DataFormat(volume.format)
+                        if volume.format is not None
+                        else DataFormat.WKW
+                    ),
                     zip=volume_path,
                     segments=segments,
                     largest_segment_id=volume.largest_segment_id,
@@ -502,6 +508,103 @@ def save(self, path: Union[str, PathLike]) -> None:
             with open(path, "wb") as f:
                 nml.write(f)
 
+    def merge_fallback_layer(
+        self,
+        target: Path,
+        dataset_directory: Path,
+        volume_layer_name: Optional[str] = None,
+        executor: Optional[Executor] = None,
+    ) -> None:
+        """
+        Merge the volume annotation with the fallback layer.
+        """
+        annotation_volumes = list(self.get_volume_layer_names())
+
+        output_dataset = Dataset(
+            target,
+            voxel_size=self.voxel_size,
+        )
+        assert (
+            len(annotation_volumes) > 0
+        ), "Annotation does not contain any volume layers!"
+
+        if volume_layer_name is not None:
+            assert (
+                volume_layer_name in annotation_volumes
+            ), f'Volume layer name "{volume_layer_name}" not found in annotation'
+        else:
+            assert (
+                len(annotation_volumes) == 1
+            ), "Volume layer name was not provided and more than one volume layer found in annotation"
+            volume_layer_name = annotation_volumes[0]
+
+        volume_layer = self._get_volume_layer(volume_layer_name=volume_layer_name)
+        fallback_layer_name = volume_layer.fallback_layer_name
+
+        if fallback_layer_name is None:
+            logging.info("No fallback layer found, save annotation as dataset.")
+            self.export_volume_layer_to_dataset(output_dataset)
+
+        else:
+            fallback_dataset_path = dataset_directory / self.dataset_name
+            fallback_layer = Dataset.open(fallback_dataset_path).get_layer(
+                fallback_layer_name
+            )
+
+            if volume_layer.zip is None:
+                logger.info("No volume annotation found. Copy fallback layer.")
+                with get_executor_for_args(args=None, executor=executor) as executor:
+                    output_dataset.add_copy_layer(
+                        fallback_layer, compress=True, executor=executor
+                    )
+
+            else:
+                tmp_annotation_layer_name = f"{self.name}-TMP"
+                logger.info(
+                    "Unpack annotation layer %s temporarily in %s as %s",
+                    volume_layer_name,
+                    output_dataset.name,
+                    tmp_annotation_layer_name,
+                )
+                # NOTE(erjel): Cannot use "temporary_volume_layer_copy" here, since tmp folders
+                # might not be accessible from slurm compute nodes.
+                input_annotation_layer = self.export_volume_layer_to_dataset(
+                    output_dataset,
+                    layer_name=tmp_annotation_layer_name,
+                    volume_layer_name=volume_layer_name,
+                )
+
+                input_annotation_mag = input_annotation_layer.get_finest_mag()
+                fallback_mag = fallback_layer.get_mag(input_annotation_mag.mag)
+
+                logger.info(
+                    "Create layer %s in %s",
+                    fallback_layer.name,
+                    output_dataset.path,
+                )
+                output_layer = output_dataset.add_layer_like(
+                    fallback_layer, fallback_layer.name
+                )
+
+                with get_executor_for_args(args=None, executor=executor) as executor:
+                    logger.info(
+                        "Copy Mag %s from %s to %s",
+                        fallback_mag.mag,
+                        fallback_layer.path,
+                        output_layer.path,
+                    )
+                    output_mag = output_layer.add_copy_mag(
+                        fallback_mag,
+                        compress=True,
+                        executor=executor,
+                    )
+
+                    output_mag.merge_with_view(input_annotation_mag, executor)
+
+                logging.info("Delete temporary annotation layer")
+                output_dataset.delete_layer(tmp_annotation_layer_name)
+                logging.info("Done.")
+
     def upload(self) -> str:
         """Uploads the annotation to your current `webknossos_context`."""
         from ..client.context import _get_api_client, _get_context
diff --git a/webknossos/webknossos/cli/merge_fallback.py b/webknossos/webknossos/cli/merge_fallback.py
--- a/webknossos/webknossos/cli/merge_fallback.py
+++ b/webknossos/webknossos/cli/merge_fallback.py
@@ -1,18 +1,13 @@
 """This module merges a volume annotation layer with its fallback layer."""
 
-import logging
 from argparse import Namespace
 from multiprocessing import cpu_count
-from typing import Any, List, Optional, Tuple
+from typing import Any, Optional
 
 import typer
-from cluster_tools import Executor
-from rich.progress import track
 from typing_extensions import Annotated
 
 from ..annotation import Annotation
-from ..dataset import Dataset, MagView
-from ..geometry import BoundingBox, Mag
 from ..utils import get_executor_for_args
 from ._utils import DistributionStrategy, parse_path
 
@@ -70,7 +65,7 @@ def main(
         ),
     ] = None,
 ) -> None:
-    """Merges a given WEBKNOSSOS annotation."""
+    """Merges a given WEBKNOSSOS annotation with a volume annotation of a fallback layer."""
 
     executor_args = Namespace(
         jobs=jobs,
@@ -78,123 +73,10 @@ def main(
         job_resources=job_resources,
     )
 
-    annotation = Annotation.load(source_annotation)
-    annotation_volumes = list(annotation.get_volume_layer_names())
-
-    output_dataset = Dataset(
-        target,
-        voxel_size=annotation.voxel_size,
-    )
-    assert len(annotation_volumes) > 0, "Annotation does not contain any volume layers!"
-
-    if volume_layer_name is not None:
-        assert (
-            volume_layer_name in annotation_volumes
-        ), f'Volume layer name "{volume_layer_name}" not found in annotation'
-    else:
-        assert (
-            len(annotation_volumes) == 1
-        ), "Volume layer name was not provided and more than one volume layer found in annotation"
-        volume_layer_name = annotation_volumes[0]
-
-    volume_layer = annotation._get_volume_layer(volume_layer_name=volume_layer_name)
-    fallback_layer_name = volume_layer.fallback_layer_name
-
-    if fallback_layer_name is None:
-        logging.info("No fallback layer found, save annotation as dataset.")
-        annotation.export_volume_layer_to_dataset(output_dataset)
-
-    else:
-        fallback_dataset_path = dataset_directory / annotation.dataset_name
-        fallback_layer = Dataset.open(fallback_dataset_path).get_layer(
-            fallback_layer_name
-        )
-
-        if volume_layer.zip is None:
-            logging.info("No volume annotation found. Copy fallback layer.")
-            with get_executor_for_args(args=executor_args) as executor:
-                output_dataset.add_copy_layer(
-                    fallback_layer, compress=True, executor=executor
-                )
-
-        else:
-            tmp_annotation_layer_name = f"{annotation.name}-TMP"
-            logging.info(
-                f"Unpack annotation layer {volume_layer_name} temporarily in {output_dataset.name} as {tmp_annotation_layer_name}"
-            )
-            # NOTE(erjel): Cannot use "temporary_volume_layer_copy" here, since tmp folders
-            # might not be accessible from slurm compute nodes.
-            input_annotation_layer = annotation.export_volume_layer_to_dataset(
-                output_dataset,
-                layer_name=tmp_annotation_layer_name,
-                volume_layer_name=volume_layer_name,
-            )
-
-            input_annotation_mag = input_annotation_layer.get_finest_mag()
-            fallback_mag = fallback_layer.get_mag(input_annotation_mag.mag)
-
-            logging.info(f"Create layer {fallback_layer.name} in {output_dataset.path}")
-            output_layer = output_dataset.add_layer_like(
-                fallback_layer, fallback_layer.name
-            )
-
-            with get_executor_for_args(args=executor_args) as executor:
-                logging.info(
-                    f"Copy Mag {fallback_mag.mag} from {fallback_layer.path} to {output_layer.path}"
-                )
-                output_mag = output_layer.add_copy_mag(
-                    fallback_mag,
-                    compress=True,
-                    executor=executor,
-                )
-
-                merge_mags(output_mag, input_annotation_mag, executor)
-
-            logging.info("Delete temporary annotation layer")
-            output_dataset.delete_layer(tmp_annotation_layer_name)
-            logging.info("Done.")
-
-
-def merge_mags(
-    output_mag: MagView,
-    input_annotation_mag: MagView,
-    executor: Executor,
-) -> None:
-    assert all(
-        input_annotation_mag.info.chunks_per_shard.to_np() == 1
-    ), "volume annotation must have file_len=1"
-    assert (
-        input_annotation_mag.info.voxel_type == output_mag.info.voxel_type
-    ), "Volume annotation must have same dtype as fallback layer"
-    assert (
-        input_annotation_mag.mag == output_mag.mag
-    ), f"Volume annotation mag {input_annotation_mag.mag} must match the fallback layer mag {output_mag.mag}"
-
-    logging.info("Scan disk for annotation shards.")
-    bboxes = list(bbox for bbox in input_annotation_mag.get_bounding_boxes_on_disk())
-
-    logging.info(f"Grouping {len(bboxes)} bboxes according to output shards.")
-    shards_with_bboxes = BoundingBox.group_boxes_with_aligned_mag(
-        bboxes, Mag(output_mag.info.shard_shape * output_mag.mag)
-    )
-
-    args = [
-        (input_annotation_mag, output_mag, shard, bboxes)
-        for shard, bboxes in shards_with_bboxes.items()
-    ]
-
-    logging.info(f"Merging {len(args)} shards.")
-    executor.map(merge_chunk, args)
-
-
-def merge_chunk(args: Tuple[MagView, MagView, BoundingBox, List[BoundingBox]]) -> None:
-    mag_in, mag_out, shard, bboxes = args
-    data_buffer = mag_out.read(absolute_bounding_box=shard)[0]
-
-    for bbox in track(bboxes, description="Processing..."):
-        read_data = mag_in.read(absolute_bounding_box=bbox)[0]
-        data_buffer[bbox.offset(-shard.topleft).in_mag(mag_in.mag).to_slices()] = (
-            read_data
+    with get_executor_for_args(args=executor_args) as executor:
+        Annotation.load(source_annotation).merge_fallback_layer(
+            target=target,
+            dataset_directory=dataset_directory,
+            volume_layer_name=volume_layer_name,
+            executor=executor,
         )
-
-    mag_out.write(data_buffer, absolute_offset=shard.topleft)
diff --git a/webknossos/webknossos/dataset/mag_view.py b/webknossos/webknossos/dataset/mag_view.py
--- a/webknossos/webknossos/dataset/mag_view.py
+++ b/webknossos/webknossos/dataset/mag_view.py
@@ -3,7 +3,7 @@
 from argparse import Namespace
 from os import PathLike
 from pathlib import Path
-from typing import TYPE_CHECKING, Iterator, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Iterator, List, Optional, Tuple, Union
 from uuid import uuid4
 
 import numpy as np
@@ -426,6 +426,52 @@ def compress(
                 True,
             )
 
+    def merge_with_view(
+        self,
+        other: "MagView",
+        executor: Executor,
+    ) -> None:
+        assert all(
+            other.info.chunks_per_shard.to_np() == 1
+        ), "volume annotation must have file_len=1"
+        assert (
+            self.info.voxel_type == other.info.voxel_type
+        ), "Volume annotation must have same dtype as fallback layer"
+        assert (
+            self.mag == other.mag
+        ), f"To merge two Views, both need the same mag: Own mag {self.mag} does not match other mag {other.mag}"
+
+        logging.info("Scan disk for annotation shards.")
+        bboxes = list(bbox for bbox in other.get_bounding_boxes_on_disk())
+
+        logging.info("Grouping %s bboxes according to output shards.", len(bboxes))
+        shards_with_bboxes = NDBoundingBox.group_boxes_with_aligned_mag(
+            bboxes, Mag(self.info.shard_shape * self.mag)
+        )
+
+        args = [(other, shard, bboxes) for shard, bboxes in shards_with_bboxes.items()]
+
+        logging.info("Merging %s shards.", len(args))
+        wait_and_ensure_success(
+            executor.map_to_futures(self.merge_chunk, args),
+            executor,
+            "Merging chunks with fallback layer",
+        )
+
+    def merge_chunk(
+        self, args: Tuple["MagView", NDBoundingBox, List[NDBoundingBox]]
+    ) -> None:
+        other, shard, bboxes = args
+        data_buffer = self.read(absolute_bounding_box=shard)[0]
+
+        for bbox in bboxes:
+            read_data = other.read(absolute_bounding_box=bbox)[0]
+            data_buffer[bbox.offset(-shard.topleft).in_mag(other.mag).to_slices()] = (
+                read_data
+            )
+
+        self.write(data_buffer, absolute_offset=shard.topleft)
+
     @property
     def _properties(self) -> MagViewProperties:
         return next(


Test patch: