Add mags of foreign datasets (#367)

* add function to add a symlink to a foreign mag

* implement add_copy_mag

* update changelog

* adjust error messages

Author: rschwanhold

Changelog.md

@@ -39,6 +39,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 
 ### Added
 - Added functions to `wkcuber.api.dataset.Dataset` and `wkcuber.api.layer.Layer` to set and get the view configuration. [#344](https://github.com/scalableminds/webknossos-cuber/pull/344)
+- Added functions to add mags of a foreign dataset (`Layer.add_symlink_mag` and `Layer.add_copy_mag`) [#367](https://github.com/scalableminds/webknossos-cuber/pull/367)
 
 ### Changed
 

tests/test_dataset.py

@@ -1150,6 +1150,81 @@ def test_add_symlink_layer() -> None:
     assert np.array_equal(original_mag.read(size=(10, 10, 10)), write_data)
 
 
+def test_add_symlink_mag(tmp_path: Path) -> None:
+    original_ds = Dataset.create(tmp_path / "original", scale=(1, 1, 1))
+    original_layer = original_ds.add_layer(
+        "color", LayerCategories.COLOR_TYPE, dtype_per_channel="uint8"
+    )
+    original_layer.add_mag(1).write(
+        data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8)
+    )
+    original_layer.add_mag(2).write(
+        data=(np.random.rand(5, 10, 15) * 255).astype(np.uint8)
+    )
+
+    ds = Dataset.create(tmp_path / "link", scale=(1, 1, 1))
+    layer = ds.add_layer("color", LayerCategories.COLOR_TYPE, dtype_per_channel="uint8")
+    layer.add_mag(1).write(
+        offset=(6, 6, 6), data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8)
+    )
+
+    assert tuple(layer.get_bounding_box().topleft) == (6, 6, 6)
+    assert tuple(layer.get_bounding_box().size) == (10, 20, 30)
+
+    symlink_mag = layer.add_symlink_mag(tmp_path / "original" / "color" / "2")
+
+    assert (tmp_path / "link" / "color" / "1").exists()
+    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 2
+
+    assert tuple(layer.get_bounding_box().topleft) == (0, 0, 0)
+    assert tuple(layer.get_bounding_box().size) == (16, 26, 36)
+
+    # Write data in symlink layer
+    # Note: The written data is fully inside the bounding box of the original data.
+    # This is important because the bounding box of the foreign layer would not be updated if we use the linked dataset to write outside of its original bounds.
+    write_data = (np.random.rand(5, 5, 5) * 255).astype(np.uint8)
+    symlink_mag.write(offset=(0, 0, 0), data=write_data)
+
+    assert np.array_equal(symlink_mag.read(size=(5, 5, 5))[0], write_data)
+    assert np.array_equal(original_layer.get_mag(2).read(size=(5, 5, 5))[0], write_data)
+
+
+def test_add_copy_mag(tmp_path: Path) -> None:
+    original_ds = Dataset.create(tmp_path / "original", scale=(1, 1, 1))
+    original_layer = original_ds.add_layer(
+        "color", LayerCategories.COLOR_TYPE, dtype_per_channel="uint8"
+    )
+    original_layer.add_mag(1).write(
+        data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8)
+    )
+    original_data = (np.random.rand(5, 10, 15) * 255).astype(np.uint8)
+    original_layer.add_mag(2).write(data=original_data)
+
+    ds = Dataset.create(tmp_path / "link", scale=(1, 1, 1))
+    layer = ds.add_layer("color", LayerCategories.COLOR_TYPE, dtype_per_channel="uint8")
+    layer.add_mag(1).write(
+        offset=(6, 6, 6), data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8)
+    )
+
+    assert tuple(layer.get_bounding_box().topleft) == (6, 6, 6)
+    assert tuple(layer.get_bounding_box().size) == (10, 20, 30)
+
+    copy_mag = layer.add_copy_mag(tmp_path / "original" / "color" / "2")
+
+    assert (tmp_path / "link" / "color" / "1").exists()
+    assert len(ds.properties.data_layers["color"].wkw_magnifications) == 2
+
+    assert tuple(layer.get_bounding_box().topleft) == (0, 0, 0)
+    assert tuple(layer.get_bounding_box().size) == (16, 26, 36)
+
+    # Write data in copied layer
+    write_data = (np.random.rand(5, 5, 5) * 255).astype(np.uint8)
+    copy_mag.write(offset=(0, 0, 0), data=write_data)
+
+    assert np.array_equal(copy_mag.read(size=(5, 5, 5))[0], write_data)
+    assert np.array_equal(original_layer.get_mag(2).read()[0], original_data)
+
+
 def test_search_dataset_also_for_long_layer_name() -> None:
     delete_dir(TESTOUTPUT_DIR / "long_layer_name")
 

wkcuber/api/layer.py

@@ -1,7 +1,9 @@
 import logging
 import math
 import os
+import shutil
 from argparse import Namespace
+from pathlib import Path
 from shutil import rmtree
 from os.path import join
 from os import makedirs
@@ -200,6 +202,88 @@ def delete_mag(self, mag: Union[int, str, list, tuple, np.ndarray, Mag]) -> None
         )
         rmtree(full_path)
 
+    def _add_foreign_mag(
+        self, foreign_mag_path: Path, symlink: bool, make_relative: bool
+    ) -> MagView:
+        mag_name = foreign_mag_path.name
+        mag = Mag(mag_name)
+        operation = "symlink" if symlink else "copy"
+        if mag in self.mags.keys():
+            raise IndexError(
+                f"Cannot {operation} {foreign_mag_path}. This dataset already has a mag called {mag_name}."
+            )
+
+        foreign_normalized_mag_path = (
+            Path(os.path.relpath(foreign_mag_path, self.dataset.path))
+            if make_relative
+            else foreign_mag_path
+        )
+
+        if symlink:
+            os.symlink(
+                foreign_normalized_mag_path,
+                join(self.dataset.path, self.name, mag_name),
+            )
+        else:
+            shutil.copytree(
+                foreign_normalized_mag_path,
+                join(self.dataset.path, self.name, mag_name),
+            )
+
+        # copy the properties of the layer into the properties of this dataset
+        mag_properties = None
+        from wkcuber.api.properties.dataset_properties import (
+            Properties,
+        )  # using a relative import prevents a circular dependency
+
+        foreign_layer_properties = Properties._from_json(
+            foreign_mag_path.parent.parent / Properties.FILE_NAME
+        ).data_layers[self.name]
+        for resolution in foreign_layer_properties.wkw_magnifications:
+            if resolution.mag == mag:
+                mag_properties = resolution
+                break
+
+        assert (
+            mag_properties is not None
+        ), f"Failed to {operation} existing mag at {foreign_mag_path}: The properties on the foreign dataset do not contain an entry for the specified mag."
+        new_bbox = self.get_bounding_box().extended_by(
+            foreign_layer_properties.get_bounding_box()
+        )
+        self.set_bounding_box(offset=new_bbox.topleft, size=new_bbox.size)
+        self.dataset.properties.data_layers[self.name]._wkw_magnifications += [
+            mag_properties
+        ]
+        self.dataset.properties._export_as_json()
+
+        self._setup_mag(mag)
+        return self.mags[mag]
+
+    def add_symlink_mag(
+        self, foreign_mag_path: Union[str, Path], make_relative: bool = False
+    ) -> MagView:
+        """
+        Creates a symlink to the data at `foreign_mag_path` which belongs to another dataset.
+        The relevant information from the `datasource-properties.json` of the other dataset is copied to this dataset.
+        Note: If the other dataset modifies its bounding box afterwards, the change does not affect this properties
+        (or vice versa).
+        If make_relative is True, the symlink is made relative to the current dataset path.
+        """
+        foreign_mag_path = Path(os.path.abspath(foreign_mag_path))
+        return self._add_foreign_mag(
+            foreign_mag_path, symlink=True, make_relative=make_relative
+        )
+
+    def add_copy_mag(self, foreign_mag_path: Union[str, Path]) -> MagView:
+        """
+        Copies the data at `foreign_mag_path` which belongs to another dataset to the current dataset.
+        Additionally, the relevant information from the `datasource-properties.json` of the other dataset are copied too.
+        """
+        foreign_mag_path = Path(os.path.abspath(foreign_mag_path))
+        return self._add_foreign_mag(
+            foreign_mag_path, symlink=False, make_relative=False
+        )
+
     def _create_dir_for_mag(
         self, mag: Union[int, str, list, tuple, np.ndarray, Mag]
     ) -> None:

wkcuber/api/properties/dataset_properties.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 from typing import Tuple, Optional, Dict, Any, cast
 
-from wkcuber.api.layer import Layer, LayerCategories
+from wkcuber.api.layer import LayerCategories
 from wkcuber.api.properties.layer_properties import (
     SegmentationLayerProperties,
     LayerProperties,