implement dask-downsampling

brainglobe_template_builder/napari/_reader.py

@@ -5,6 +5,7 @@
 implement multiple readers or even other plugin contributions. see:
 https://napari.org/stable/plugins/guides.html?#readers
 """
+
 from pathlib import Path
 
 from brainglobe_utils.IO.image.load import load_any

brainglobe_template_builder/preproc/transform_utils.py

@@ -1,5 +1,7 @@
+import dask.array as da
 import numpy as np
 from scipy.ndimage import affine_transform
+from skimage import transform
 
 
 def get_rotation_from_vectors(vec1: np.ndarray, vec2: np.ndarray):
@@ -71,3 +73,60 @@ def apply_transform(
     # Preserve original data range and type
     transformed = np.clip(transformed, data.min(), data.max())
     return transformed.astype(data.dtype)
+
+
+def downsample_anisotropic_image_stack(
+    stack: da.Array, xy_downsampling: int, z_downsampling: int
+) -> np.ndarray:
+    """
+
+    Lazily downsamples a dask array first along axis 1,2 and then along axis 0,
+    using a local mean of the pixels. The (smaller) array is returned
+    in memory (numpy) form at the end.
+
+    This setup is typical for certain types of microscopy,
+    where z-resolution is lower than x-y-resolution.
+
+    The input dask array must be chunked by x-y slice,
+
+    Parameters:
+    ----------
+    stack : da.Array
+        The input dask array representing the image stack.
+    xy_downsampling : int
+        The downsampling factor for the x and y axes.
+    z_downsampling : int
+        The downsampling factor for the z axis.
+    Returns:
+    -------
+    np.ndarray
+        The computed downsampled (numpy) array.
+    Raises:
+    ------
+    AssertionError
+        If the array is not chunked slice-wise on axis 0.
+    """
+    # check we have expected slice chunks
+    assert np.all(
+        np.array(stack.chunks[0]) == 1
+    ), f"Array not chunked slice-wise! Chunks on axis 0 are {stack.chunks[0]}"
+
+    # we have xy slices as chunks, so apply downscaling in xy first
+    downsampled_xy = stack.map_blocks(
+        transform.downscale_local_mean,
+        (1, xy_downsampling, xy_downsampling),
+        dtype=np.float64,
+    )
+
+    # rechunk so we can map_blocks along z
+    downsampled_xy = downsampled_xy.rechunk(
+        {0: downsampled_xy.shape[0], 1: -1, 2: -1}
+    )
+
+    # downsample in z
+    downsampled_z = downsampled_xy.map_blocks(
+        transform.downscale_local_mean,
+        (z_downsampling, 1, 1),
+        dtype=np.float64,
+    )
+    return downsampled_z.compute()

tests/test_unit/test_transform_utils.py

@@ -0,0 +1,50 @@
+import dask.array as da
+import numpy as np
+import pytest
+from skimage import transform
+
+from brainglobe_template_builder.preproc.transform_utils import (
+    downsample_anisotropic_image_stack,
+)
+
+
+@pytest.fixture()
+def stack():
+    """Create a dask array representing an image stack"""
+    data = np.random.rand(10, 100, 100)  # Random image stack
+    return da.from_array(data, chunks=(1, 100, 100))
+
+
+@pytest.fixture()
+def not_slicewise_stack():
+    """Create a dask array representing an image stack"""
+    data = np.random.rand(10, 100, 100)  # Random image stack
+    return da.from_array(data, chunks=(2, 100, 100))
+
+
+def test_downsample_anisotropic_image_stack(stack):
+    """Test that downsampling with dask gives same as without."""
+    xy_downsampling = 20
+    z_downsampling = 2
+
+    downsampled_stack = downsample_anisotropic_image_stack(
+        stack, xy_downsampling, z_downsampling
+    )
+
+    assert downsampled_stack.shape == (5, 5, 5)
+
+    expected = transform.downscale_local_mean(
+        stack.compute(), (1, xy_downsampling, xy_downsampling)
+    )
+    expected = transform.downscale_local_mean(expected, (z_downsampling, 1, 1))
+    assert np.all(
+        downsampled_stack == expected
+    ), "dask downsampling does not match expected skimage result"
+
+
+def test_downsample_anisotropic_image_stack_raises(not_slicewise_stack):
+    with pytest.raises(AssertionError) as e:
+        downsample_anisotropic_image_stack(
+            not_slicewise_stack, xy_downsampling=20, z_downsampling=2
+        )
+    assert e.match("not chunked slice-wise!")