PERF: itkwasm downsampling in zyx order

Support memory layout assumed in _large_image_serialization.

Author: thewtex

ngff_zarr/methods/_itkwasm.py

@@ -12,7 +12,7 @@
     _dim_scale_factors,
     _get_block,
     _spatial_dims,
-    _spatial_dims_last,
+    _spatial_dims_last_zyx,
 )
 
 _image_dims: Tuple[str, str, str, str] = ("x", "y", "z", "t")
@@ -96,7 +96,7 @@ def _downsample_itkwasm(
         previous_dim_factors = dim_factors
         previous_image = _align_chunks(previous_image, default_chunks, dim_factors)
         # Operate on a contiguous spatial block
-        previous_image = _spatial_dims_last(previous_image)
+        previous_image = _spatial_dims_last_zyx(previous_image)
         if previous_image.dims != dims:
             transposed_dims = True
             reorder = [previous_image.dims.index(dim) for dim in dims]
@@ -270,7 +270,7 @@ def _downsample_itkwasm(
                 out_chunks_list.append(1)
         downscaled_array = downscaled_array.rechunk(tuple(out_chunks_list))
 
-        # transpose back to original order if needed (_spatial_dims_last transposed the order)
+        # transpose back to original order if needed (_spatial_dims_zyx transposed the order)
         # breakpoint()
         if transposed_dims:
             downscaled_array = downscaled_array.transpose(reorder)

ngff_zarr/methods/_support.py

@@ -42,6 +42,44 @@ def _spatial_dims_last(ngff_image: NgffImage) -> NgffImage:
     return result
 
 
+def _spatial_dims_last_zyx(ngff_image: NgffImage) -> NgffImage:
+    dims = list(ngff_image.dims)
+    spatial_dims = [dim for dim in dims if dim in _spatial_dims]
+
+    # If spatial dimensions are already zyx, return the original image
+    if spatial_dims == ["z", "y", "x"] or spatial_dims == ["y", "x"]:
+        dims_spatial_channel = len(spatial_dims)
+        if dims[-1] == "c":
+            dims_spatial_channel += 1
+
+        # If spatial dimensions are already last (and 'c' can be last), return the original image
+        if all(dim in dims[-dims_spatial_channel:] for dim in spatial_dims + ["c"]):
+            return ngff_image
+
+    # Move spatial dimensions to the end, keeping 'c' as the last pre-spatial dimension if present
+    non_spatial_dims = [dim for dim in dims if dim not in _spatial_dims]
+    new_spatial_dims = ["z", "y", "x"][-len(spatial_dims) :]
+    if "c" in non_spatial_dims:
+        non_spatial_dims.remove("c")
+        new_dims = non_spatial_dims + ["c"] + new_spatial_dims
+    else:
+        new_dims = non_spatial_dims + new_spatial_dims
+
+    new_order = [dims.index(dim) for dim in new_dims]
+
+    if tuple(new_dims) == tuple(ngff_image.dims):
+        return ngff_image
+
+    # Reorder the data array
+    reordered_data = ngff_image.data.transpose(new_order)
+
+    result = copy.copy(ngff_image)
+    result.data = reordered_data
+    result.dims = tuple(new_dims)
+
+    return result
+
+
 def _channel_dim_last(ngff_image: NgffImage) -> NgffImage:
     if "c" not in ngff_image.dims or ngff_image.dims[-1] == "c":
         return ngff_image