Merge pull request #21 from brainglobe/dask_downsampling

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,63 @@ 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, in_plane_factor: int, axial_factor: int
+) -> np.ndarray:
+    """
+
+    Lazily downsamples a dask array first along axes 1,2 (in-plane) and then
+    along axis 0 (axial), using a local mean of the pixels. The image is
+    zero-padded to allow for the correct dimensions of the averaging
+    neighbourhood, since it uses `skimage.transform.downscale_local_mean`
+    under the hood.
+
+    This setup is typical for certain types of microscopy,
+    where axial resolution is lower than in-plane resolution.
+
+    The input dask array must be chunked by plane. The (smaller) array
+    is returned in memory (numpy) form at the end.
+
+    Parameters:
+    ----------
+    stack : da.Array
+        The input dask array representing the image stack.
+    in_plane_factor : int
+        The in-plane downsampling factor (axes=1,2).
+    axial_factor : int
+        The downsampling factor in axial direction (axis=0).
+    Returns:
+    -------
+    np.ndarray
+        The computed downsampled (numpy) array.
+    Raises:
+    ------
+    AssertionError
+        If the array is not chunked by plane along axis 0.
+    """
+    # check we have expected slice chunks
+    assert np.all(
+        np.array(stack.chunks[0]) == 1
+    ), f"Array not chunked by plane! Chunks on axis 0 are {stack.chunks[0]}"
+
+    # we have xy slices as chunks, so apply downscaling in xy first
+    downsampled_inplane = stack.map_blocks(
+        transform.downscale_local_mean,
+        (1, in_plane_factor, in_plane_factor),
+        dtype=np.float64,
+    )
+
+    # rechunk so we can map_blocks along z
+    downsampled_inplane = downsampled_inplane.rechunk(
+        {0: downsampled_inplane.shape[0], 1: -1, 2: -1}
+    )
+
+    # downsample in z
+    downsampled_axial = downsampled_inplane.map_blocks(
+        transform.downscale_local_mean,
+        (axial_factor, 1, 1),
+        dtype=np.float64,
+    )
+    return downsampled_axial.compute()

tests/test_unit/test_transform_utils.py

@@ -0,0 +1,49 @@
+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"""
+    rng = np.random.default_rng()
+    data = rng.random(size=(10, 100, 100))
+    return da.from_array(data, chunks=(1, 100, 100))
+
+
+@pytest.fixture()
+def not_slicewise_stack(stack):
+    """Create a dask array representing an image stack"""
+    return stack.rechunk({0: 2, 1: 100, 2: 100})
+
+
+def test_downsample_anisotropic_image_stack(stack):
+    """Test that downsampling with dask gives same as without."""
+    in_plane = 20
+    axial = 2
+
+    downsampled_stack = downsample_anisotropic_image_stack(
+        stack, in_plane, axial
+    )
+
+    assert downsampled_stack.shape == (5, 5, 5)
+
+    expected = transform.downscale_local_mean(
+        stack.compute(), (1, in_plane, in_plane)
+    )
+    expected = transform.downscale_local_mean(expected, (axial, 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, match="not chunked by plane!"):
+        downsample_anisotropic_image_stack(
+            not_slicewise_stack, in_plane_factor=20, axial_factor=2
+        )