Add shallow_copy_dataset method (#437)

* add shallow_copy_dataset and test

* add test for make_relative and fix make_relative bugs

* fix typing

* correct comments

* fix lint

* add_layer/mag_for_existing_files

* use add_mag_for_existing_files

* implement feedback

* correct comments

* add changelog

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

Author: valentin-pinkau

webknossos/Changelog.md

@@ -11,6 +11,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 
 ### Breaking Changes
 ### Added
+- The Dataset class now has a new method: add_shallow_copy. [#437](https://github.com/scalableminds/webknossos-libs/pull/437)
 ### Changed
 - The `Vec3Int` constructor now asserts that its components are whole numbers also in numpy case. [#434](https://github.com/scalableminds/webknossos-libs/pull/434)
 - Updated scikit-image dependency to 0.18.3. [#435](https://github.com/scalableminds/webknossos-libs/pull/435)

webknossos/tests/test_dataset.py

@@ -1232,9 +1232,10 @@ def test_add_symlink_mag(tmp_path: Path) -> None:
     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)
-    )
+    original_mag_2 = original_layer.add_mag(2)
+    original_mag_2.write(data=(np.random.rand(5, 10, 15) * 255).astype(np.uint8))
+    original_mag_4 = original_layer.add_mag(4)
+    original_mag_4.write(data=(np.random.rand(2, 5, 7) * 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")
@@ -1245,10 +1246,11 @@ def test_add_symlink_mag(tmp_path: Path) -> None:
     assert tuple(layer.bounding_box.topleft) == (6, 6, 6)
     assert tuple(layer.bounding_box.size) == (10, 20, 30)
 
-    symlink_mag = layer.add_symlink_mag(tmp_path / "original" / "color" / "2")
+    symlink_mag_2 = layer.add_symlink_mag(original_mag_2)
+    symlink_mag_4 = layer.add_symlink_mag(original_mag_4.path)
 
     assert (tmp_path / "link" / "color" / "1").exists()
-    assert len(layer._properties.wkw_resolutions) == 2
+    assert len(layer._properties.wkw_resolutions) == 3
 
     assert tuple(layer.bounding_box.topleft) == (0, 0, 0)
     assert tuple(layer.bounding_box.size) == (16, 26, 36)
@@ -1257,9 +1259,9 @@ def test_add_symlink_mag(tmp_path: Path) -> None:
     # 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)
+    symlink_mag_2.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(symlink_mag_2.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)
 
     assure_exported_properties(ds)
@@ -1275,7 +1277,8 @@ def test_add_copy_mag(tmp_path: Path) -> None:
         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)
+    original_mag_2 = original_layer.add_mag(2)
+    original_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")
@@ -1286,7 +1289,7 @@ def test_add_copy_mag(tmp_path: Path) -> None:
     assert tuple(layer.bounding_box.topleft) == (6, 6, 6)
     assert tuple(layer.bounding_box.size) == (10, 20, 30)
 
-    copy_mag = layer.add_copy_mag(tmp_path / "original" / "color" / "2")
+    copy_mag = layer.add_copy_mag(original_mag_2)
 
     assert (tmp_path / "link" / "color" / "1").exists()
     assert len(layer._properties.wkw_resolutions) == 2
@@ -1359,6 +1362,46 @@ def test_outdated_dtype_parameter() -> None:
         )
 
 
+@pytest.mark.parametrize("make_relative", [True, False])
+def test_dataset_shallow_copy(make_relative: bool) -> None:
+    delete_dir(TESTOUTPUT_DIR / "original_dataset")
+    delete_dir(TESTOUTPUT_DIR / "copy_dataset")
+    ds = Dataset.create(TESTOUTPUT_DIR / "original_dataset", (1, 1, 1))
+    original_layer_1 = ds.add_layer(
+        "color", LayerCategories.COLOR_TYPE, dtype_per_layer=np.uint8, num_channels=1
+    )
+    original_layer_1.add_mag(1)
+    original_layer_1.add_mag("2-2-1")
+    original_layer_2 = ds.add_layer(
+        "segmentation",
+        LayerCategories.SEGMENTATION_TYPE,
+        dtype_per_layer=np.uint32,
+        largest_segment_id=0,
+    )
+    original_layer_2.add_mag(4)
+    mappings_path = original_layer_2.path / "mappings"
+    os.makedirs(mappings_path)
+    open(mappings_path / "agglomerate_view.hdf5", "w").close()
+
+    shallow_copy_of_ds = ds.shallow_copy_dataset(
+        TESTOUTPUT_DIR / "copy_dataset", make_relative=make_relative
+    )
+    shallow_copy_of_ds.get_layer("color").add_mag(Mag("4-4-1"))
+    assert (
+        len(Dataset(TESTOUTPUT_DIR / "original_dataset").get_layer("color").mags) == 2
+    ), "Adding a new mag should not affect the original dataset"
+    assert (
+        len(Dataset(TESTOUTPUT_DIR / "copy_dataset").get_layer("color").mags) == 3
+    ), "Expecting all mags from original dataset and new downsampled mag"
+    assert os.path.exists(
+        TESTOUTPUT_DIR
+        / "copy_dataset"
+        / "segmentation"
+        / "mappings"
+        / "agglomerate_view.hdf5"
+    ), "Expecting mappings to exist in shallow copy"
+
+
 def test_dataset_conversion() -> None:
     origin_ds_path = TESTOUTPUT_DIR / "conversion" / "origin_wk"
     converted_ds_path = TESTOUTPUT_DIR / "conversion" / "converted_wk"

webknossos/webknossos/dataset/dataset.py

@@ -14,7 +14,7 @@
 import wkw
 
 from webknossos.geometry import BoundingBox, Vec3Int
-from webknossos.utils import get_executor_for_args
+from webknossos.utils import copy_directory_with_symlinks, get_executor_for_args
 
 from .layer import (
     Layer,
@@ -348,6 +348,27 @@ def add_layer_like(self, other_layer: Layer, layer_name: str) -> Layer:
         self._export_as_json()
         return self._layers[layer_name]
 
+    def add_layer_for_existing_files(
+        self, layer_name: str, category: str, **kwargs: Any
+    ) -> Layer:
+        assert layer_name not in self.layers, f"Layer {layer_name} already exists!"
+        mag_headers = list((self.path / layer_name).glob("*/header.wkw"))
+        assert (
+            len(mag_headers) != 0
+        ), f"Could not find any header.wkw files in {self.path / layer_name}, cannot add layer."
+        with wkw.Dataset.open(str(mag_headers[0].parent)) as wkw_dataset:
+            header = wkw_dataset.header
+        layer = self.add_layer(
+            layer_name,
+            category=category,
+            num_channels=header.num_channels,
+            dtype_per_channel=header.voxel_type,
+            **kwargs,
+        )
+        for mag_dir in layer.path.iterdir():
+            layer.add_mag_for_existing_files(mag_dir.name)
+        return layer
+
     def get_segmentation_layer(self) -> SegmentationLayer:
         """
         Returns the only segmentation layer.
@@ -535,6 +556,36 @@ def copy_dataset(
         new_ds._export_as_json()
         return new_ds
 
+    def shallow_copy_dataset(
+        self,
+        new_dataset_path: Path,
+        name: Optional[str] = None,
+        make_relative: bool = False,
+    ) -> "Dataset":
+        """
+        Create a new dataset at the given path. Link all mags of all existing layers.
+        In addition, link all other directories in all layer directories
+        to make this method robust against additional files e.g. layer/mappings/agglomerate_view.hdf5.
+        This method becomes useful when exposing a dataset to webknossos.
+        """
+        new_dataset = Dataset.create(
+            new_dataset_path, scale=self.scale, name=name or self.name
+        )
+        for layer_name, layer in self.layers.items():
+            new_layer = new_dataset.add_layer_like(layer, layer_name)
+            for mag_view in layer.mags.values():
+                new_layer.add_symlink_mag(mag_view, make_relative)
+
+            # copy all other directories with a dir scan
+            copy_directory_with_symlinks(
+                layer.path,
+                new_layer.path,
+                ignore=[str(mag) for mag in layer.mags] + [PROPERTIES_FILE_NAME],
+                make_relative=make_relative,
+            )
+
+        return new_dataset
+
     def _get_layer_by_category(self, category: str) -> Layer:
         assert (
             category == LayerCategories.COLOR_TYPE

webknossos/webknossos/dataset/layer.py

@@ -1,4 +1,3 @@
-import copy
 import logging
 import math
 import operator
@@ -310,6 +309,29 @@ def add_mag(
 
         return self._mags[mag]
 
+    def add_mag_for_existing_files(
+        self,
+        mag: Union[int, str, list, tuple, np.ndarray, Mag],
+    ) -> MagView:
+        """
+        Creates a new mag based on already existing files.
+
+        Raises an IndexError if the specified `mag` does not exists.
+        """
+        mag = Mag(mag)
+        assert (
+            mag not in self.mags
+        ), f"Cannot add mag {mag} as it already exists for layer {self}"
+        self._setup_mag(mag)
+        mag_view = self._mags[mag]
+        cube_length = mag_view.header.file_len * mag_view.header.block_len
+        self._properties.wkw_resolutions.append(
+            MagViewProperties(resolution=mag, cube_length=cube_length)
+        )
+        self.dataset._export_as_json()
+
+        return mag_view
+
     def get_or_add_mag(
         self,
         mag: Union[int, str, list, tuple, np.ndarray, Mag],
@@ -372,71 +394,93 @@ 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
+        self,
+        foreign_mag_view_or_path: Union[os.PathLike, str, MagView],
+        symlink: bool,
+        make_relative: bool,
+        extend_layer_bounding_box: bool = True,
     ) -> 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}."
+        """
+        The foreign mag is (shallow) copied and the existing mag is added to the datasource-properties.json.
+        If extend_layer_bounding_box is true, the self.bounding_box will be extended
+        by the bounding box of the layer the foreign mag belongs to.
+        """
+
+        if isinstance(foreign_mag_view_or_path, MagView):
+            foreign_mag_view = foreign_mag_view_or_path
+        else:
+            # local import to prevent circular dependency
+            from .dataset import Dataset
+
+            foreign_mag_view_path = Path(foreign_mag_view_or_path)
+            foreign_mag_view = (
+                Dataset(foreign_mag_view_path.parent.parent)
+                .get_layer(foreign_mag_view_path.parent.name)
+                .get_mag(foreign_mag_view_path.name)
             )
 
+        self._assert_mag_does_not_exist_yet(foreign_mag_view.mag)
+
         foreign_normalized_mag_path = (
-            Path(os.path.relpath(foreign_mag_path, self.dataset.path))
+            Path(os.path.relpath(foreign_mag_view.path, self.path))
             if make_relative
-            else foreign_mag_path
+            else Path(os.path.abspath(foreign_mag_view.path))
         )
 
         if symlink:
             os.symlink(
                 foreign_normalized_mag_path,
-                join(self.dataset.path, self.name, mag_name),
+                join(self.dataset.path, self.name, str(foreign_mag_view.mag)),
             )
         else:
             shutil.copytree(
                 foreign_normalized_mag_path,
-                join(self.dataset.path, self.name, mag_name),
+                join(self.dataset.path, self.name, str(foreign_mag_view.mag)),
             )
 
-        # copy the properties of the layer into the properties of this dataset
-        from .dataset import Dataset  # local import to prevent circular dependency
-
-        original_layer = Dataset(foreign_mag_path.parent.parent).get_layer(
-            foreign_mag_path.parent.name
-        )
-        original_mag = original_layer.get_mag(foreign_mag_path.name)
-        mag_properties = copy.deepcopy(original_mag._properties)
-
-        self.bounding_box = self.bounding_box.extended_by(original_layer.bounding_box)
-        self._properties.wkw_resolutions += [mag_properties]
-        self._setup_mag(mag)
+        self.add_mag_for_existing_files(foreign_mag_view.mag)
+        if extend_layer_bounding_box:
+            self.bounding_box = self.bounding_box.extended_by(
+                foreign_mag_view.layer.bounding_box
+            )
         self.dataset._export_as_json()
-        return self.mags[mag]
+
+        return self._mags[foreign_mag_view.mag]
 
     def add_symlink_mag(
-        self, foreign_mag_path: Union[str, Path], make_relative: bool = False
+        self,
+        foreign_mag_view_or_path: Union[os.PathLike, str, MagView],
+        make_relative: bool = False,
+        extend_layer_bounding_box: bool = True,
     ) -> MagView:
         """
-        Creates a symlink to the data at `foreign_mag_path` which belongs to another dataset.
+        Creates a symlink to the data at `foreign_mag_view_or_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
+            foreign_mag_view_or_path,
+            symlink=True,
+            make_relative=make_relative,
+            extend_layer_bounding_box=extend_layer_bounding_box,
         )
 
-    def add_copy_mag(self, foreign_mag_path: Union[str, Path]) -> MagView:
+    def add_copy_mag(
+        self,
+        foreign_mag_view_or_path: Union[os.PathLike, str, MagView],
+        extend_layer_bounding_box: bool = True,
+    ) -> MagView:
         """
-        Copies the data at `foreign_mag_path` which belongs to another dataset to the current dataset.
+        Copies the data at `foreign_mag_view_or_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
+            foreign_mag_view_or_path,
+            symlink=False,
+            make_relative=False,
+            extend_layer_bounding_box=extend_layer_bounding_box,
         )
 
     def _create_dir_for_mag(

webknossos/webknossos/dataset/mag_view.py

@@ -2,7 +2,6 @@
 import os
 import shutil
 from argparse import Namespace
-from os.path import join
 from pathlib import Path
 from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union, cast
 from uuid import uuid4
@@ -94,14 +93,16 @@ def __init__(
         )
 
         if create:
-            wkw.Dataset.create(
-                join(layer.dataset.path, layer.name, self.name), self.header
-            )
+            wkw.Dataset.create(str(self.path), self.header)
 
     @property
     def layer(self) -> "Layer":
         return self._layer
 
+    @property
+    def path(self) -> Path:
+        return self._path
+
     @property
     def _properties(self) -> MagViewProperties:
         return next(