WIP: tczyz support

Author: thewtex

ngff_zarr/methods/_itkwasm.py

@@ -10,6 +10,7 @@
     _dim_scale_factors,
     _get_block,
     _spatial_dims,
+    _spatial_dims_last,
 )
 
 _image_dims: Tuple[str, str, str, str] = ("x", "y", "z", "t")
@@ -89,13 +90,14 @@ def _downsample_itkwasm_bin_shrink(
         dim_factors = _dim_scale_factors(dims, scale_factor, previous_dim_factors)
         previous_dim_factors = dim_factors
         previous_image = _align_chunks(previous_image, default_chunks, dim_factors)
+        previous_image = _spatial_dims_last(previous_image)
 
         shrink_factors = [dim_factors[sd] for sd in spatial_dims]
 
         block_0 = _get_block(previous_image, 0)
 
         # For consistency for now, do not utilize direction until there is standardized support for
-        # direction cosines / orientation in OME-NGFF
+        # direction cosines / orientation in OME-NGFF (v0.6)
         # block_0.attrs.pop("direction", None)
         block_input = itkwasm.image_from_array(np.ones_like(block_0))
         spacing = [previous_image.scale[d] for d in spatial_dims]
@@ -166,6 +168,7 @@ def _downsample_itkwasm(
         dim_factors = _dim_scale_factors(dims, scale_factor, previous_dim_factors)
         previous_dim_factors = dim_factors
         previous_image = _align_chunks(previous_image, default_chunks, dim_factors)
+        previous_image = _spatial_dims_last(previous_image)
 
         shrink_factors = [dim_factors[sd] for sd in spatial_dims]
 
@@ -177,7 +180,10 @@ def _downsample_itkwasm(
         block_neg1_input = _get_block(previous_image, -1)
 
         # Compute overlap for Gaussian blurring for all blocks
-        block_0_image = itkwasm.image_from_array(np.ones_like(block_0_input))
+        is_vector = previous_image.dims[-1] == "c"
+        block_0_image = itkwasm.image_from_array(
+            np.ones_like(block_0_input), is_vector=is_vector
+        )
         input_spacing = [previous_image.scale[d] for d in spatial_dims]
         block_0_image.spacing = input_spacing
         input_origin = [previous_image.translation[d] for d in spatial_dims]
@@ -208,17 +214,25 @@ def _downsample_itkwasm(
             block_output.size[dim] == computed_size[dim]
             for dim in range(block_output.data.ndim)
         )
+        breakpoint()
         output_chunks = list(previous_image.data.chunks)
-        if "t" in previous_image.dims:
-            dims = list(previous_image.dims)
-            t_index = dims.index("t")
-            output_chunks.pop(t_index)
+        dims = list(previous_image.dims)
+        output_chunks_start = 0
+        while dims[output_chunks_start] not in _spatial_dims:
+            output_chunks_start += 1
+        output_chunks = output_chunks[output_chunks_start:]
+        # if "t" in previous_image.dims:
+        #     dims = list(previous_image.dims)
+        #     t_index = dims.index("t")
+        #     output_chunks.pop(t_index)
         for i, c in enumerate(output_chunks):
             output_chunks[i] = [
                 block_output.data.shape[i],
             ] * len(c)
         # Compute output size for block N-1
-        block_neg1_image = itkwasm.image_from_array(np.ones_like(block_neg1_input))
+        block_neg1_image = itkwasm.image_from_array(
+            np.ones_like(block_neg1_input), is_vector=is_vector
+        )
         block_neg1_image.spacing = input_spacing
         block_neg1_image.origin = input_origin
         block_output = downsample_bin_shrink(

ngff_zarr/methods/_support.py

@@ -1,17 +1,75 @@
 from typing import List
+import copy
 
-from dask.array import take
 
 from ..ngff_image import NgffImage
 
 _spatial_dims = {"x", "y", "z"}
 
+_spatial_dims = {"x", "y", "z"}
+
+
+def _spatial_dims_last(ngff_image: NgffImage) -> NgffImage:
+    dims = list(ngff_image.dims)
+    spatial_dims = [dim for dim in dims if dim in _spatial_dims]
+
+    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]
+    if "c" in non_spatial_dims:
+        non_spatial_dims.remove("c")
+        new_dims = non_spatial_dims + ["c"] + spatial_dims
+    else:
+        new_dims = non_spatial_dims + 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
+
+    dims = list(ngff_image.dims)
+    # Move 'c' dimension to the end
+    dims.remove("c")
+    dims.append("c")
+
+    # Reorder the data array
+    new_order = [ngff_image.dims.index(dim) for dim in dims]
+    reordered_data = ngff_image.data.transpose(new_order)
+
+    result = copy.copy(ngff_image)
+    result.data = reordered_data
+    result.dims = tuple(dims)
+
+    return result
+
 
 def _dim_scale_factors(dims, scale_factor, previous_dim_factors):
     if isinstance(scale_factor, int):
         result_scale_factors = {
             dim: int(scale_factor / previous_dim_factors[dim])
-            for dim in dims if dim in _spatial_dims
+            for dim in dims
+            if dim in _spatial_dims
         }
     else:
         result_scale_factors = {
@@ -71,11 +129,15 @@ def _get_block(previous_image: NgffImage, block_index: int):
     """Helper method for accessing an enumerated chunk from input"""
     block_shape = [c[block_index] for c in previous_image.data.chunks]
     block = previous_image.data[tuple([slice(0, s) for s in block_shape])]
-    # For consistency for now, do not utilize direction until there is standardized support for
-    # direction cosines / orientation in OME-NGFF
-    # block.attrs.pop("direction", None)
-    if "t" in previous_image.dims:
-        dims = list(previous_image.dims)
-        t_index = dims.index("t")
-        block = take(block, 0, t_index)
+    dims = list(previous_image.dims)
+    # Also take "c" if it is the last dimension
+    if dims[-1] == "c":
+        dims[-1] = "x"
+    indexer = []
+    for d in dims:
+        if d in _spatial_dims:
+            indexer.append(slice(None))
+        else:
+            indexer.append(0)
+    block = block[tuple(indexer)]
     return block

ngff_zarr/ngff_image_to_itk_image.py

@@ -2,6 +2,7 @@
 
 from .methods._support import _spatial_dims
 from .ngff_image import NgffImage
+from .methods._support import _channel_dim_last
 
 
 def _dtype_to_component_type(dtype):
@@ -37,6 +38,8 @@ def ngff_image_to_itk_image(
 ):
     from itkwasm import IntTypes, PixelTypes
 
+    ngff_image = _channel_dim_last(ngff_image)
+
     dims = ngff_image.dims
     dimension = 3 if "z" in dims else 2
 
@@ -63,7 +66,6 @@ def ngff_image_to_itk_image(
     origin = [ngff_image.translation[dim] for dim in spatial_dims]
     size = [ngff_image.data.shape[dims.index(d)] for d in spatial_dims]
 
-    # TODO: reorder as needed
     data = np.asarray(ngff_image.data)
 
     image_dict = {

test/test_to_ngff_zarr_itkwasm.py

@@ -1,4 +1,7 @@
-from ngff_zarr import Methods, to_multiscales
+import numpy as np
+from zarr.storage import MemoryStore
+
+from ngff_zarr import Methods, to_multiscales, to_ngff_image, to_ngff_zarr
 
 from ._data import verify_against_baseline
 
@@ -11,6 +14,14 @@
     pass
 
 
+def test_channel_support():
+    data = np.random.randint(0, 256, 16777216).reshape((2, 128, 256, 256))
+    image = to_ngff_image(data, dims=["c", "z", "y", "x"])
+    multiscales = to_multiscales(image, scale_factors=[2, 4], chunks=64)
+    store = MemoryStore()
+    to_ngff_zarr(store, multiscales)
+
+
 def test_bin_shrink_isotropic_scale_factors(input_images):
     dataset_name = "cthead1"
     image = input_images[dataset_name]