add read_cxyz and write_cxyz

webknossos/Changelog.md

@@ -13,8 +13,10 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v3.4.3...HEAD)
 
 ### Breaking Changes
+- `View.read_xyz` is deprecated in favor of `View.read_cxyz`. [#1461](https://github.com/scalableminds/webknossos-libs/pull/1461)
 
 ### Added
+- Added `View.read_cxyz` and `View.write_cxyz` to read and write arrays with guaranteed cxyz axis ordering. [#1461](https://github.com/scalableminds/webknossos-libs/pull/1461)
 
 ### Changed
 

webknossos/pyproject.toml

@@ -71,13 +71,14 @@ Changelog = "https://github.com/scalableminds/webknossos-libs/blob/master/webkno
 
 # A list of all of the optional dependencies. They can be opted into by other apps.
 [project.optional-dependencies]
+fastremap = ["fastremap ~=1.15.0"]
 tifffile = ["tifffile >=2024.8.24"]
 imagecodecs = ["imagecodecs >=2021.11.20"]
 bioformats = ["JPype1 ~=1.5.0"]
 czi = ["pylibCZIrw ==5.1.1"]
 s3 = ["s3fs >=2026.2.0"]
 examples = [
-    "fastremap ~=1.15.0",
+    "webknossos[fastremap]",
     "pandas ~=2.2.0",
     "pooch ~=1.5.2",
     "tabulate >=0.9.0",
@@ -86,6 +87,7 @@ examples = [
 ]
 all = [
     "webknossos[tifffile]",
+    "webknossos[fastremap]",
     "webknossos[imagecodecs]",
     "webknossos[bioformats]",
     "webknossos[czi]",

webknossos/tests/dataset/test_dataset.py

@@ -42,7 +42,7 @@
     SegmentationLayerProperties,
 )
 from webknossos.dataset_properties.structuring import get_dataset_converter
-from webknossos.geometry import BoundingBox, Mag, Vec3Int, VecIntLike
+from webknossos.geometry import BoundingBox, Mag, NDBoundingBox, Vec3Int, VecIntLike
 from webknossos.utils import (
     copytree,
     dump_path,
@@ -606,6 +606,159 @@ def test_view_write(data_format: DataFormat, output_path: UPath) -> None:
     np.testing.assert_array_equal(data, write_data)
 
 
+@pytest.mark.parametrize("data_format,output_path", DATA_FORMATS_AND_OUTPUT_PATHS)
+def test_read_cxyz(data_format: DataFormat, output_path: UPath) -> None:
+    ds_path = copy_simple_dataset(data_format, output_path)
+
+    with pytest.warns(UserWarning, match=".*not aligned with the shard shape.*"):
+        wk_view = (
+            Dataset.open(ds_path)
+            .get_layer("color")
+            .get_mag("1")
+            .get_view(absolute_offset=(0, 0, 0), size=(16, 16, 16))
+        )
+
+    data_cxyz = wk_view.read_cxyz(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    assert data_cxyz.shape == (3, 10, 10, 10)
+
+    # read_cxyz must return the same data as read() for standard axis ordering
+    data = wk_view.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    np.testing.assert_array_equal(data_cxyz, data)
+
+
+@pytest.mark.parametrize("data_format,output_path", DATA_FORMATS_AND_OUTPUT_PATHS)
+def test_write_cxyz(data_format: DataFormat, output_path: UPath) -> None:
+    ds_path = copy_simple_dataset(data_format, output_path)
+    with pytest.warns(UserWarning, match=".*not aligned with the shard shape.*"):
+        wk_view = (
+            Dataset.open(ds_path)
+            .get_layer("color")
+            .get_mag("1")
+            .get_view(absolute_offset=(0, 0, 0), size=(16, 16, 16))
+        )
+
+    np.random.seed(1234)
+    write_data = (np.random.rand(3, 10, 10, 10) * 255).astype(np.uint8)
+
+    wk_view.write_cxyz(write_data, allow_unaligned=True)
+
+    data = wk_view.read_cxyz(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    np.testing.assert_array_equal(data, write_data)
+
+
+@pytest.mark.parametrize("data_format,output_path", DATA_FORMATS_AND_OUTPUT_PATHS)
+def test_read_cxyz_adds_channel_axis(
+    data_format: DataFormat, output_path: UPath
+) -> None:
+    ds_path = prepare_dataset_path(data_format, output_path)
+    layer = Dataset(ds_path, voxel_size=(1, 1, 1)).add_layer(
+        "segmentation",
+        SEGMENTATION_CATEGORY,
+        bounding_box=NDBoundingBox((0, 0, 0), (10, 10, 10), axes="xyz"),
+        data_format=data_format,
+        num_channels=1,
+    )
+    mag = layer.add_mag("1")
+
+    write_data = np.zeros((10, 10, 10), dtype=np.uint64)
+    mag.write(write_data, absolute_offset=(0, 0, 0))
+
+    data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    assert data.shape == (10, 10, 10)
+
+    data = mag.read_cxyz(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    assert data.shape == (1, 10, 10, 10)
+
+
+@pytest.mark.parametrize(
+    "layer_bbox,write_bbox,write_data,expected_shape",
+    [
+        (
+            NDBoundingBox(topleft=(0, 0), size=(10, 20), axes=("x", "y"), index=(0, 1)),
+            NDBoundingBox(topleft=(0, 0), size=(10, 20), axes=("x", "y"), index=(0, 1)),
+            np.arange(200, dtype=np.uint8).reshape(1, 10, 20, 1),
+            (1, 10, 20, 1),
+        ),
+        (
+            NDBoundingBox(
+                topleft=(0, 0, 0, 0, 0),
+                size=(1, 4, 4, 4, 2),
+                axes=("c", "x", "y", "z", "t"),
+                index=(0, 1, 2, 3, 4),
+            ),
+            NDBoundingBox(
+                topleft=(0, 0, 0, 0, 0),
+                size=(1, 4, 4, 4, 1),
+                axes=("c", "x", "y", "z", "t"),
+                index=(0, 1, 2, 3, 4),
+            ),
+            np.zeros((1, 4, 4, 4), dtype=np.uint8),
+            (1, 4, 4, 4),
+        ),
+        (
+            NDBoundingBox(
+                topleft=(0, 0, 0),
+                size=(10, 20, 5),
+                axes=("x", "y", "z"),
+                index=(0, 1, 2),
+            ),
+            NDBoundingBox(
+                topleft=(0, 0, 0),
+                size=(10, 20, 5),
+                axes=("x", "y", "z"),
+                index=(0, 1, 2),
+            ),
+            (np.arange(1000, dtype=np.uint8)).reshape(1, 10, 20, 5),
+            (1, 10, 20, 5),
+        ),
+    ],
+)
+def test_read_write_cxyz_axes(
+    tmp_path: UPath,
+    layer_bbox: NDBoundingBox,
+    write_bbox: NDBoundingBox,
+    write_data: np.ndarray,
+    expected_shape: tuple[int, ...],
+) -> None:
+    """read_cxyz/write_cxyz must work when the bounding box has no z axis."""
+    mag = (
+        Dataset(tmp_path / "ds", voxel_size=(1, 1, 1))
+        .add_layer("color", COLOR_CATEGORY, bounding_box=layer_bbox)
+        .add_mag("1")
+    )
+
+    mag.write_cxyz(write_data, absolute_bounding_box=write_bbox)
+
+    data = mag.read_cxyz(absolute_bounding_box=write_bbox)
+    assert data.shape == expected_shape, f"unexpected shape {data.shape}"
+    np.testing.assert_array_equal(data, write_data)
+
+
+@pytest.mark.parametrize("data_format,output_path", DATA_FORMATS_AND_OUTPUT_PATHS)
+def test_write_cxyz_mag_view(data_format: DataFormat, output_path: UPath) -> None:
+    ds_path = prepare_dataset_path(data_format, output_path)
+    layer = Dataset(ds_path, voxel_size=(1, 1, 1)).add_layer(
+        "color", COLOR_CATEGORY, num_channels=3
+    )
+    mag = layer.add_mag("1")
+
+    np.random.seed(1234)
+    write_data = (np.random.rand(3, 10, 10, 10) * 255).astype(np.uint8)
+
+    # without allow_resize should fail
+    with pytest.raises(
+        ValueError, match=".*does not fit in the layer's bounding box.*"
+    ):
+        mag.write_cxyz(write_data, absolute_offset=(0, 0, 0))
+
+    # with allow_resize should succeed and update the bounding box
+    mag.write_cxyz(write_data, absolute_offset=(0, 0, 0), allow_resize=True)
+
+    assert layer.bounding_box == BoundingBox((0, 0, 0), (10, 10, 10))
+    data = mag.read_cxyz(absolute_offset=(0, 0, 0), size=(10, 10, 10))
+    np.testing.assert_array_equal(data, write_data)
+
+
 @pytest.mark.parametrize("output_path", [TESTOUTPUT_DIR, REMOTE_TESTOUTPUT_DIR])
 @pytest.mark.parametrize("data_format", [DataFormat.Zarr, DataFormat.Zarr3])
 def test_direct_zarr_access(output_path: UPath, data_format: DataFormat) -> None:
@@ -1431,7 +1584,11 @@ def test_changing_layer_bounding_box(
     assert original_data.shape == (3, 24, 24, 24)
 
     layer.bounding_box = layer.bounding_box.with_size(
-        [12, 12, 10]
+        [
+            12,
+            12,
+            10,
+        ]
     )  # decrease bounding box
 
     bbox_size = ds.get_layer("color").bounding_box.size
@@ -1441,7 +1598,11 @@ def test_changing_layer_bounding_box(
     np.testing.assert_array_equal(original_data[:, :12, :12, :10], less_data)
 
     layer.bounding_box = layer.bounding_box.with_size(
-        [36, 48, 60]
+        [
+            36,
+            48,
+            60,
+        ]
     )  # increase the bounding box
 
     bbox_size = ds.get_layer("color").bounding_box.size

webknossos/webknossos/cli/export_as_tiff.py

@@ -5,7 +5,6 @@
 from math import ceil
 from typing import Annotated, Any
 
-import fastremap
 import numpy as np
 import typer
 from cluster_tools import Executor
@@ -71,6 +70,8 @@ def _slice_to_image(data_slice: np.ndarray, downsample: int = 1) -> np.ndarray:
 
 
 def _apply_mapping(data: np.ndarray, mapping_array: TensorStore) -> np.ndarray:
+    import fastremap
+
     unique_ids = fastremap.unique(data)
     in_bounds = [i for i in unique_ids if i < mapping_array.shape[0]]
     out_of_bounds = [i for i in unique_ids if i >= mapping_array.shape[0]]
@@ -300,6 +301,14 @@ def main(
         mag_view = layer.get_mag(mag)
 
         if apply_mapping is not None:
+            try:
+                import fastremap  # noqa: F401
+            except ImportError:
+                raise ImportError(
+                    "fastremap is required to use --apply-mapping. "
+                    "Please install with `pip install webknossos[fastremap]`."
+                )
+
             if layer.category != SEGMENTATION_CATEGORY:
                 raise ValueError(
                     f"--apply-mapping requires a segmentation layer, "

webknossos/webknossos/client/_download_dataset.py

@@ -128,5 +128,5 @@ def download_dataset(
                 data = np.frombuffer(chunk_bytes, dtype=layer.dtype).reshape(
                     layer.num_channels, *chunk_in_mag.size, order="F"
                 )
-                mag_view.write(data, absolute_offset=chunk.topleft)
+                mag_view.write_cxyz(data, absolute_offset=chunk.topleft)
     return dataset

webknossos/webknossos/dataset/_utils/pims_czi_reader.py

@@ -9,7 +9,7 @@
     from pylibCZIrw import czi as pyczi
 except ImportError as e:
     raise ImportError(
-        "Cannot import pylibCZIrw, please install it e.g. using pip install --extra-index-url https://pypi.scm.io/simple/ webknossos[czi]"
+        "Cannot import pylibCZIrw, please install it e.g. using pip install webknossos[czi]"
     ) from e
 
 PIXEL_TYPE_TO_DTYPE = {

webknossos/webknossos/dataset/layer/_downsampling_utils.py

@@ -321,7 +321,7 @@ def downsample_cube_job(
 
             bbox = source_bbox.offset(source_offset).with_size_xyz(source_size)
 
-            cube_buffer_channels = source_view.read_xyz(
+            cube_buffer_channels = source_view.read_cxyz(
                 absolute_bounding_box=bbox,
             )
 

webknossos/webknossos/dataset/layer/_upsampling_utils.py

@@ -52,7 +52,7 @@ def upsample_cube_job(
         ):
             shape = chunk.in_mag(target_view.mag).size.to_tuple()
             file_buffer = np.zeros(shape, dtype=target_view.get_dtype())
-            cube_buffer_channels = source_view.read_xyz(
+            cube_buffer_channels = source_view.read_cxyz(
                 absolute_bounding_box=chunk,
             )
 

webknossos/webknossos/dataset/layer/view/mag_view.py

@@ -415,6 +415,54 @@ def write(
             absolute_bounding_box=mag1_bbox,
         )
 
+    def write_cxyz(
+        self,
+        data: np.ndarray,
+        *,
+        allow_resize: bool = False,
+        allow_unaligned: bool = False,
+        relative_offset: Vec3IntLike | None = None,  # in mag1
+        absolute_offset: Vec3IntLike | None = None,  # in mag1
+        relative_bounding_box: NDBoundingBox | None = None,  # in mag1
+        absolute_bounding_box: NDBoundingBox | None = None,  # in mag1
+    ) -> None:
+        """Write data from a (c, x, y, z) ordered array to the magnification level.
+
+        Equivalent to :meth:`write` but always accepts data in ``(c, x, y, z)`` axis
+        order regardless of the underlying storage axis ordering.
+
+        Args:
+            data (np.ndarray): 4D array in ``(c, x, y, z)`` order.
+            allow_resize (bool, optional): If True, allows updating the layer's
+                bounding box if the write extends beyond it. Defaults to False.
+            allow_unaligned (bool, optional): If True, allows writing data without
+                being aligned to the shard shape. Defaults to False.
+            relative_offset (Vec3IntLike | None, optional): Offset relative to view's
+                position in Mag(1) coordinates. Defaults to None.
+            absolute_offset (Vec3IntLike | None, optional): Absolute offset in Mag(1)
+                coordinates. Defaults to None.
+            relative_bounding_box (NDBoundingBox | None, optional): Bounding box
+                relative to view's position in Mag(1) coordinates. Defaults to None.
+            absolute_bounding_box (NDBoundingBox | None, optional): Absolute bounding
+                box in Mag(1) coordinates. Defaults to None.
+        """
+        assert len(data.shape) == 4, (
+            f"write_cxyz expects a 4D (c, x, y, z) array, got shape {data.shape}"
+        )
+        data, write_loc = self._resolve_cxyz_write(
+            data,
+            relative_offset,
+            absolute_offset,
+            relative_bounding_box,
+            absolute_bounding_box,
+        )
+        self.write(
+            data,
+            allow_resize=allow_resize,
+            allow_unaligned=allow_unaligned,
+            **write_loc,
+        )
+
     def get_bounding_boxes_on_disk(
         self,
     ) -> Iterator[NDBoundingBox]: