diff --git a/brainglobe_template_builder/napari/_reader.py b/brainglobe_template_builder/napari/_reader.py index 114015b..19fb080 100644 --- a/brainglobe_template_builder/napari/_reader.py +++ b/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 diff --git a/brainglobe_template_builder/preproc/transform_utils.py b/brainglobe_template_builder/preproc/transform_utils.py index fc4738a..fcd0185 100644 --- a/brainglobe_template_builder/preproc/transform_utils.py +++ b/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() diff --git a/tests/test_unit/test_transform_utils.py b/tests/test_unit/test_transform_utils.py new file mode 100644 index 0000000..0785d76 --- /dev/null +++ b/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 + )