Fix upsampling (#1287)

* wip fixing upsample ds.

* wip replace offset size combinations with bboxes.

* rewrite with chunk method.

* lint and typecheck.

* remove assertion.

* fix tests.

* Update readme and fix tests.

* Update webknossos/Changelog.md

---------

Co-authored-by: Daniel <[email protected]>

Author: markbader

webknossos/Changelog.md

@@ -13,13 +13,15 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v2.1.0...HEAD)
 
 ### Breaking Changes
+- If `buffer_shape` is passed to `Layer.upsample()` it must be a multiple of the `shard_shape`. [#1287](https://github.com/scalableminds/webknossos-libs/pull/1287)
 
 ### Added
 
 ### Changed
 - Using `x-auth-header` to send tokens to datastore. [#1270](https://github.com/scalableminds/webknossos-libs/pull/1270)
 
 ### Fixed
+- Fixed an issue with upsampling views. [#1287](https://github.com/scalableminds/webknossos-libs/pull/1287)
 
 
 ## [2.1.0](https://github.com/scalableminds/webknossos-libs/releases/tag/v2.1.0) - 2025-04-01

webknossos/tests/dataset/test_upsampling.py

@@ -42,7 +42,6 @@ def test_upsampling(tmp_path: Path) -> None:
         finest_mag=Mag(1),
         compress=False,
         sampling_mode=SamplingModes.ANISOTROPIC,
-        buffer_shape=(64, 64, 64),
     )
 
     assert layer.get_mag("2").read().mean() == layer.get_mag("1").read().mean()
@@ -92,7 +91,7 @@ def upsample_test_helper(tmp_path: Path, use_compress: bool) -> None:
             0,
         ),
         [0.5, 0.5, 1.0],
-        BUFFER_SHAPE,
+        mag1.info.shard_shape,
     )
 
     assert np.any(source_buffer != 0)
@@ -138,7 +137,7 @@ def test_upsample_multi_channel(tmp_path: Path) -> None:
     upsample_cube_job(
         (mag2.get_view(), layer.get_mag("1").get_view(), 0),
         [0.5, 0.5, 0.5],
-        BUFFER_SHAPE,
+        mag2.info.shard_shape,
     )
 
     channels = [

webknossos/webknossos/dataset/_downsampling_utils.py

@@ -28,11 +28,8 @@ class InterpolationModes(Enum):
     MIN = 6
 
 
-DEFAULT_BUFFER_SHAPE = Vec3Int.full(256)
-
-
 def determine_buffer_shape(array_info: ArrayInfo) -> Vec3Int:
-    return DEFAULT_BUFFER_SHAPE.pairmin(array_info.shard_shape)
+    return array_info.shard_shape
 
 
 def calculate_mags_to_downsample(

webknossos/webknossos/dataset/_upsampling_utils.py

@@ -1,10 +1,9 @@
 import logging
-import math
-from itertools import product
 
 import numpy as np
 
 from ..geometry import Vec3Int
+from .data_format import DataFormat
 from .view import View
 
 
@@ -32,55 +31,42 @@ def upsample_cube_job(
     assert all(1 >= f for f in mag_factors), (
         f"mag_factors ({mag_factors}) for upsampling must be smaller than 1"
     )
+    if (
+        target_view._data_format != DataFormat.WKW
+        and target_view._is_compressed() == False
+    ):
+        assert buffer_shape % target_view.info.shard_shape == Vec3Int.zeros(), (
+            f"buffer_shape ({buffer_shape}) must be divisible by shard_shape ({target_view.info.shard_shape})"
+        )
+    inverse_factors = [int(1 / f) for f in mag_factors]
 
     try:
         num_channels = target_view.info.num_channels
-        target_bbox_in_mag = target_view.bounding_box.in_mag(target_view.mag)
-        shape = (num_channels,) + target_bbox_in_mag.size.to_tuple()
-        shape_xyz = target_bbox_in_mag.size_xyz
-        file_buffer = np.zeros(shape, target_view.get_dtype())
-
-        tiles = product(
-            *list(
-                [
-                    list(range(0, math.ceil(length)))
-                    for length in shape_xyz.to_np() / buffer_shape.to_np()
-                ]
-            )
-        )
 
-        for tile in tiles:
-            target_offset = Vec3Int(tile) * buffer_shape
-            source_offset = target_offset * source_view.mag
-            source_size = source_view.bounding_box.size_xyz
-            source_size = (buffer_shape * source_view.mag).pairmin(
-                source_size - source_offset
-            )
-
-            bbox = source_view.bounding_box.offset(source_offset).with_size_xyz(
-                source_size
-            )
+        for chunk in target_view.bounding_box.chunk(
+            buffer_shape * target_view.mag, buffer_shape * target_view.mag
+        ):
+            shape = (num_channels,) + (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(
-                absolute_bounding_box=bbox,
+                absolute_bounding_box=chunk,
             )
 
             for channel_index in range(num_channels):
                 cube_buffer = cube_buffer_channels[channel_index]
 
                 if not np.all(cube_buffer == 0):
-                    # Upsample the buffer
-                    inverse_factors = [int(1 / f) for f in mag_factors]
                     data_cube = upsample_cube(cube_buffer, inverse_factors)
 
-                    buffer_bbox = target_view.bounding_box.with_topleft_xyz(
-                        target_offset * inverse_factors
-                    ).with_size_xyz(data_cube.shape)
+                    buffer_bbox = chunk.with_topleft_xyz(Vec3Int.zeros()).with_size_xyz(
+                        data_cube.shape
+                    )
                     data_cube = buffer_bbox.xyz_array_to_bbox_shape(data_cube)
                     file_buffer[(channel_index,) + buffer_bbox.to_slices_xyz()] = (
                         data_cube
                     )
 
-        target_view.write(file_buffer, absolute_bounding_box=target_view.bounding_box)
+            target_view.write(file_buffer, absolute_bounding_box=chunk)
 
     except Exception as exc:
         logging.error(