Fixes the array shape when using Dataset.from_images (#1267)

* fixes the Zarr array shape for Dataset.from_images

* changelog

* rm expand_only

* docstring

Author: normanrz

webknossos/Changelog.md

@@ -21,6 +21,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 
 ### Fixed
 - Fixed an issue with `RemoteDataset.explore_and_add_remote()` where including a remote dataset failed. [#1261](https://github.com/scalableminds/webknossos-libs/pull/1261)
+- Fixed an issue with the Zarr array's shape when using `Dataset.from_images`. [#1267](https://github.com/scalableminds/webknossos-libs/pull/1267)
 
 
 ## [2.0.0](https://github.com/scalableminds/webknossos-libs/releases/tag/v2.0.0) - 2025-03-04

webknossos/tests/dataset/test_from_images.py

@@ -1,3 +1,4 @@
+import json
 import warnings
 from pathlib import Path
 from shutil import copytree
@@ -8,7 +9,6 @@
 from cluster_tools import SequentialExecutor
 from tifffile import TiffFile
 
-import webknossos as wk
 from tests.constants import TESTDATA_DIR
 from webknossos.dataset import Dataset
 
@@ -21,15 +21,15 @@ def ignore_warnings() -> Iterator:
 
 
 def test_compare_tifffile(tmp_path: Path) -> None:
-    ds = wk.Dataset.from_images(
+    ds = Dataset.from_images(
         TESTDATA_DIR / "tiff",
         tmp_path,
         (1, 1, 1),
         compress=True,
         layer_name="tiff_stack",
         layer_category="segmentation",
         shard_shape=(256, 256, 256),
-        map_filepath_to_layer_name=wk.Dataset.ConversionLayerMapping.ENFORCE_SINGLE_LAYER,
+        map_filepath_to_layer_name=Dataset.ConversionLayerMapping.ENFORCE_SINGLE_LAYER,
     )
     assert len(ds.layers) == 1
     assert "tiff_stack" in ds.layers
@@ -41,7 +41,7 @@ def test_compare_tifffile(tmp_path: Path) -> None:
 
 
 def test_multiple_multitiffs(tmp_path: Path) -> None:
-    ds = wk.Dataset.from_images(
+    ds = Dataset.from_images(
         TESTDATA_DIR / "various_tiff_formats",
         tmp_path,
         (1, 1, 1),
@@ -71,9 +71,18 @@ def test_multiple_multitiffs(tmp_path: Path) -> None:
         assert layer.num_channels == channels
         assert layer.bounding_box.size == size
 
+        # Check that the zarr.json metadata is correct
+        mag1 = layer.get_finest_mag()
+        array_shape = json.loads((mag1.path / "zarr.json").read_bytes())["shape"]
+        shard_aligned_bottomright = layer.bounding_box.with_bottomright_xyz(
+            layer.bounding_box.bottomright_xyz.ceildiv(mag1.info.shard_shape)
+            * mag1.info.shard_shape
+        ).bottomright
+        assert array_shape == [channels] + shard_aligned_bottomright.to_list()
+
 
 def test_from_dicom_images(tmp_path: Path) -> None:
-    ds = wk.Dataset.from_images(
+    ds = Dataset.from_images(
         TESTDATA_DIR / "dicoms",
         tmp_path,
         (1, 1, 1),
@@ -97,7 +106,7 @@ def test_no_slashes_in_layername(tmp_path: Path) -> None:
 
     for strategy in Dataset.ConversionLayerMapping:
         with SequentialExecutor() as executor:
-            dataset = wk.Dataset.from_images(
+            dataset = Dataset.from_images(
                 tmp_path / "tiff",
                 tmp_path / str(strategy),
                 voxel_size=(10, 10, 10),

webknossos/webknossos/dataset/_array.py

@@ -1,5 +1,4 @@
 import re
-import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from dataclasses import dataclass
@@ -99,12 +98,7 @@ def write(self, bbox: NDBoundingBox, data: np.ndarray) -> None:
         pass
 
     @abstractmethod
-    def ensure_size(
-        self,
-        new_bbox: NDBoundingBox,
-        align_with_shards: bool = True,
-        warn: bool = False,
-    ) -> None:
+    def resize(self, new_bbox: NDBoundingBox) -> None:
         pass
 
     @abstractmethod
@@ -206,12 +200,7 @@ def read(self, bbox: NDBoundingBox) -> np.ndarray:
     def write(self, bbox: NDBoundingBox, data: np.ndarray) -> None:
         self._wkw_dataset.write(Vec3Int(bbox.topleft), data)
 
-    def ensure_size(
-        self,
-        new_bbox: NDBoundingBox,
-        align_with_shards: bool = True,
-        warn: bool = False,
-    ) -> None:
+    def resize(self, new_bbox: NDBoundingBox) -> None:
         pass
 
     def _list_files(self) -> Iterator[Path]:
@@ -511,42 +500,22 @@ def read(self, bbox: NDBoundingBox) -> np.ndarray:
             out = np.expand_dims(out, 0)
         return out
 
-    def ensure_size(
-        self,
-        new_bbox: NDBoundingBox,
-        align_with_shards: bool = True,
-        warn: bool = False,
-    ) -> None:
+    def resize(self, new_bbox: NDBoundingBox) -> None:
         array = self._array
 
-        new_bbox = new_bbox.with_bottomright(
-            (
-                max(array.domain.exclusive_max[i + 1], new_bbox.bottomright[i])
-                for i in range(len(new_bbox))
-            )
+        # Align with shards
+        shard_shape = self.info.shard_shape
+        new_bbox = new_bbox.with_bottomright_xyz(
+            new_bbox.bottomright_xyz.ceildiv(shard_shape) * shard_shape
         )
         new_domain = tensorstore.IndexDomain(
             new_bbox.ndim + 1,
             shape=(self.info.num_channels,) + new_bbox.bottomright.to_tuple(),
             implicit_upper_bounds=tuple(True for _ in range(new_bbox.ndim + 1)),
             labels=array.domain.labels,
         )
-        if new_domain != array.domain:
-            if align_with_shards:
-                shard_shape = self.info.shard_shape
-                new_aligned_bbox = new_bbox.with_bottomright_xyz(
-                    new_bbox.bottomright_xyz.ceildiv(shard_shape) * shard_shape
-                )
-                new_domain = tensorstore.IndexDomain(
-                    new_aligned_bbox.ndim + 1,
-                    shape=(self.info.num_channels,)
-                    + new_aligned_bbox.bottomright.to_tuple(),
-                    implicit_upper_bounds=tuple(
-                        True for _ in range(new_aligned_bbox.ndim + 1)
-                    ),
-                    labels=array.domain.labels,
-                )
 
+        if new_domain != array.domain:
             # Check on-disk for changes to shape
             current_array = tensorstore.open(
                 {
@@ -560,17 +529,10 @@ def ensure_size(
                     + "This is likely happening because multiple processes changed the metadata of this array."
                 )
 
-            if warn:
-                warnings.warn(
-                    f"[INFO] Resizing Zarr array from `{array.domain}` to `{new_domain}`.",
-                    category=UserWarning,
-                )
-
             self._cached_array = array.resize(
                 inclusive_min=None,
                 exclusive_max=new_domain.exclusive_max,
                 resize_metadata_only=True,
-                expand_only=True,
             ).result()
 
     def write(self, bbox: NDBoundingBox, data: np.ndarray) -> None:

webknossos/webknossos/dataset/layer.py

@@ -357,15 +357,13 @@ def bounding_box(self) -> NDBoundingBox:
 
     @bounding_box.setter
     def bounding_box(self, bbox: NDBoundingBox) -> None:
-        """
-        Updates the offset and size of the bounding box of this layer in the properties.
-        """
+        """Updates the offset and size of the bounding box of this layer in the properties."""
         self.dataset._ensure_writable()
         assert bbox.topleft.is_positive(), f"Updating the bounding box of layer {self} to {bbox} failed, topleft must not contain negative dimensions."
         self._properties.bounding_box = bbox
         self.dataset._export_as_json()
         for mag in self.mags.values():
-            mag._array.ensure_size(bbox.align_with_mag(mag.mag).in_mag(mag.mag))
+            mag._array.resize(bbox.align_with_mag(mag.mag).in_mag(mag.mag))
 
     @property
     def category(self) -> LayerCategoryType:
@@ -561,7 +559,7 @@ def add_mag(
             path=mag_path,
         )
 
-        mag_view._array.ensure_size(
+        mag_view._array.resize(
             self.bounding_box.align_with_mag(mag, ceil=True).in_mag(mag)
         )