From 174dcaa85b1ff4ec7b3943dca8c43b3c8e6efb7c Mon Sep 17 00:00:00 2001 From: alessandrofelder Date: Thu, 15 Aug 2024 11:59:18 +0100 Subject: [PATCH 1/4] implement dask-downsampling --- .../preproc/transform_utils.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/brainglobe_template_builder/preproc/transform_utils.py b/brainglobe_template_builder/preproc/transform_utils.py index fc4738a..361b7b4 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,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() From b7bd0cf24f43e46825757e443cb75ad15a2075df Mon Sep 17 00:00:00 2001 From: alessandrofelder Date: Thu, 15 Aug 2024 11:59:35 +0100 Subject: [PATCH 2/4] implement dask-downsampling test --- tests/test_unit/test_transform_utils.py | 50 +++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 tests/test_unit/test_transform_utils.py diff --git a/tests/test_unit/test_transform_utils.py b/tests/test_unit/test_transform_utils.py new file mode 100644 index 0000000..e7718c8 --- /dev/null +++ b/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!") From 7cb67c328c11bddea2d700b8b3df51ceb8b83233 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Aug 2024 11:01:15 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- brainglobe_template_builder/napari/_reader.py | 1 + 1 file changed, 1 insertion(+) 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 From 6e12335d9147be6c57e181b83e38b539cb3f89ca Mon Sep 17 00:00:00 2001 From: alessandrofelder Date: Thu, 15 Aug 2024 14:47:09 +0100 Subject: [PATCH 4/4] improve docs, variable names, test --- .../preproc/transform_utils.py | 41 ++++++++++--------- tests/test_unit/test_transform_utils.py | 23 +++++------ 2 files changed, 33 insertions(+), 31 deletions(-) diff --git a/brainglobe_template_builder/preproc/transform_utils.py b/brainglobe_template_builder/preproc/transform_utils.py index 361b7b4..fcd0185 100644 --- a/brainglobe_template_builder/preproc/transform_utils.py +++ b/brainglobe_template_builder/preproc/transform_utils.py @@ -76,27 +76,30 @@ def apply_transform( def downsample_anisotropic_image_stack( - stack: da.Array, xy_downsampling: int, z_downsampling: int + stack: da.Array, in_plane_factor: int, axial_factor: 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. + 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 z-resolution is lower than x-y-resolution. + where axial resolution is lower than in-plane resolution. - The input dask array must be chunked by x-y slice, + 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. - xy_downsampling : int - The downsampling factor for the x and y axes. - z_downsampling : int - The downsampling factor for the z axis. + 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 @@ -104,29 +107,29 @@ def downsample_anisotropic_image_stack( Raises: ------ AssertionError - If the array is not chunked slice-wise on axis 0. + 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 slice-wise! Chunks on axis 0 are {stack.chunks[0]}" + ), 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_xy = stack.map_blocks( + downsampled_inplane = stack.map_blocks( transform.downscale_local_mean, - (1, xy_downsampling, xy_downsampling), + (1, in_plane_factor, in_plane_factor), 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} + downsampled_inplane = downsampled_inplane.rechunk( + {0: downsampled_inplane.shape[0], 1: -1, 2: -1} ) # downsample in z - downsampled_z = downsampled_xy.map_blocks( + downsampled_axial = downsampled_inplane.map_blocks( transform.downscale_local_mean, - (z_downsampling, 1, 1), + (axial_factor, 1, 1), dtype=np.float64, ) - return downsampled_z.compute() + return downsampled_axial.compute() diff --git a/tests/test_unit/test_transform_utils.py b/tests/test_unit/test_transform_utils.py index e7718c8..0785d76 100644 --- a/tests/test_unit/test_transform_utils.py +++ b/tests/test_unit/test_transform_utils.py @@ -11,40 +11,39 @@ @pytest.fixture() def stack(): """Create a dask array representing an image stack""" - data = np.random.rand(10, 100, 100) # Random 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(): +def not_slicewise_stack(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)) + 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.""" - xy_downsampling = 20 - z_downsampling = 2 + in_plane = 20 + axial = 2 downsampled_stack = downsample_anisotropic_image_stack( - stack, xy_downsampling, z_downsampling + stack, in_plane, axial ) assert downsampled_stack.shape == (5, 5, 5) expected = transform.downscale_local_mean( - stack.compute(), (1, xy_downsampling, xy_downsampling) + stack.compute(), (1, in_plane, in_plane) ) - expected = transform.downscale_local_mean(expected, (z_downsampling, 1, 1)) + 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) as e: + with pytest.raises(AssertionError, match="not chunked by plane!"): downsample_anisotropic_image_stack( - not_slicewise_stack, xy_downsampling=20, z_downsampling=2 + not_slicewise_stack, in_plane_factor=20, axial_factor=2 ) - assert e.match("not chunked slice-wise!")