diff --git a/brainglobe_utils/image/binning.py b/brainglobe_utils/image/binning.py index 383b2ac..ba015fd 100644 --- a/brainglobe_utils/image/binning.py +++ b/brainglobe_utils/image/binning.py @@ -11,5 +11,5 @@ def get_bins(image_size, bin_sizes): """ bins = [] for dim in range(0, len(image_size)): - bins.append(np.arange(0, image_size[dim], bin_sizes[dim])) + bins.append(np.arange(0, image_size[dim] + 1, bin_sizes[dim])) return bins diff --git a/brainglobe_utils/image/heatmap.py b/brainglobe_utils/image/heatmap.py new file mode 100644 index 0000000..9c306b5 --- /dev/null +++ b/brainglobe_utils/image/heatmap.py @@ -0,0 +1,109 @@ +from pathlib import Path +from typing import Tuple, Union + +import imio +import numpy as np +from scipy.ndimage import zoom +from skimage.filters import gaussian + +from brainglobe_utils.general.system import ensure_directory_exists +from brainglobe_utils.image.binning import get_bins +from brainglobe_utils.image.masking import mask_image_threshold +from brainglobe_utils.image.scale import scale_and_convert_to_16_bits + + +def rescale_array(source_array, target_array, order=1): + """ + Rescale a source array to the size of a target array. + + Parameters + ---------- + source_array : np.ndarray + The array to be rescaled. + target_array : np.ndarray + The array whose size will be matched. + order : int (default=1) + Returns + ------- + np.ndarray + The rescaled source array. + """ + + # Compute the zoom factors for each dimension + zoom_factors = [ + t / s for s, t in zip(source_array.shape, target_array.shape) + ] + + # Use scipy's zoom function to rescale the array + rescaled_array = zoom(source_array, zoom_factors, order=order) + + return rescaled_array + + +def heatmap_from_points( + points: np.ndarray, + image_resolution: float, + image_shape: Tuple[int, int, int], + bin_sizes: Tuple[int, int, int] = (1, 1, 1), + output_filename: Union[str, Path, None] = None, + smoothing: Union[float, None] = None, + mask_image: Union[np.ndarray, None] = None, +) -> np.ndarray: + """ + Generate a heatmap from a set of points and a reference image that + defines the shape and resolution of the heatmap. + + Parameters + ---------- + points : np.ndarray + Data points to be used for heatmap generation. + These may be cell positions, e.g. from cellfinder. + image_resolution : float + Resolution of the image to be generated. + image_shape : Tuple[int, int, int] + The shape of the downsampled heatmap as a tuple of integers + bin_sizes : Tuple[int, int, int], optional + The bin sizes for each dimension of the heatmap (default is (1, 1, 1)). + output_filename : Union[str, pathlib.Path, None] + The filename or path where the heatmap image will be saved. + Can be a string or a pathlib.Path object. If None, the heatmap will + not be saved + smoothing : float, optional + Smoothing factor to be applied to the heatmap. Expressed in "real" + units (e.g. microns). This will be converted to voxel units based on + the image_resolution. + If None, no smoothing is applied (default is None). + mask_image : Union[np.ndarray, None], optional + An optional mask image (as a NumPy array) to apply to the heatmap. + This could be e.g. a registered atlas image (so points outside the + brain are masked). The image will be + binarised, so this image can be binary or not. + If None, no mask is applied (default is None). + + Returns + ------- + np.ndarray + The generated heatmap as a NumPy array. + """ + + bins = get_bins(image_shape, bin_sizes) + heatmap_array, _ = np.histogramdd(points, bins=bins) + heatmap_array = heatmap_array.astype(np.uint16) + + if smoothing is not None: + smoothing = int(round(smoothing / image_resolution)) + heatmap_array = gaussian(heatmap_array, sigma=smoothing) + + if mask_image is not None: + if mask_image.shape != heatmap_array.shape: + mask_image = rescale_array(mask_image, heatmap_array) + + heatmap_array = mask_image_threshold(heatmap_array, mask_image) + + heatmap_array = scale_and_convert_to_16_bits(heatmap_array) + + if output_filename is not None: + ensure_directory_exists(Path(output_filename).parent) + imio.to_tiff(heatmap_array, output_filename) + + return heatmap_array diff --git a/pyproject.toml b/pyproject.toml index af7f542..15a1583 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,9 @@ dependencies = [ "PyYAML", "scipy", "numpy", - "slurmio" + "slurmio", + "tifffile", + "imio" ] license = {text = "MIT"} diff --git a/tests/data/image/heatmap.tif b/tests/data/image/heatmap.tif new file mode 100644 index 0000000..31a8d9e Binary files /dev/null and b/tests/data/image/heatmap.tif differ diff --git a/tests/tests/test_image/test_binning.py b/tests/tests/test_image/test_binning.py index c5896ec..aad5442 100644 --- a/tests/tests/test_image/test_binning.py +++ b/tests/tests/test_image/test_binning.py @@ -7,9 +7,9 @@ def test_get_bins(): image_size = (1000, 20, 125, 725) bin_sizes = (500, 2, 25, 100) - dim0_bins = np.array((0, 500)) - dim1_bins = np.array((0, 2, 4, 6, 8, 10, 12, 14, 16, 18)) - dim2_bins = np.array((0, 25, 50, 75, 100)) + dim0_bins = np.array((0, 500, 1000)) + dim1_bins = np.array((0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20)) + dim2_bins = np.array((0, 25, 50, 75, 100, 125)) dim3_bins = np.array((0, 100, 200, 300, 400, 500, 600, 700)) bins = binning.get_bins(image_size, bin_sizes) diff --git a/tests/tests/test_image/test_heatmap.py b/tests/tests/test_image/test_heatmap.py new file mode 100644 index 0000000..8b91bf1 --- /dev/null +++ b/tests/tests/test_image/test_heatmap.py @@ -0,0 +1,48 @@ +from pathlib import Path + +import numpy as np +import pytest +from tifffile import imread + +from brainglobe_utils.image.heatmap import heatmap_from_points, rescale_array + +data_dir = Path("tests", "data") +heatmap_validate_path = data_dir / "image" / "heatmap.tif" + +points = np.array([[5, 5, 5], [10, 10, 10], [15, 15, 15]]) + + +@pytest.fixture +def mask_array(): + """ + A simple array of size (20,20,20) where the left side + is 0 and the right side is 1 + """ + array = np.zeros((20, 20, 20)) + array[:, :, 10:] = 1 + return array + + +def test_rescale_array(mask_array): + small_array = np.zeros((10, 10, 10)) + resized_array = rescale_array(mask_array, small_array) + assert resized_array.shape == small_array.shape + + +def test_heatmap_from_points(tmp_path, mask_array): + output_filename = tmp_path / "test_heatmap.tif" + + heatmap_test = heatmap_from_points( + points, + 2, + (20, 20, 20), + bin_sizes=(2, 2, 2), + output_filename=output_filename, + smoothing=5, + mask_image=mask_array, + ) + heatmap_validate = imread(heatmap_validate_path) + assert np.array_equal(heatmap_validate, heatmap_test) + + heatmap_file = imread(output_filename) + assert np.array_equal(heatmap_validate, heatmap_file)