Brainmapper utils (#22)

* copy binning bugfix from brainmapper

* Add heatmap functionality from brainmapper

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update dependencies

* Update path to test data

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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

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

pyproject.toml

@@ -15,7 +15,9 @@ dependencies = [
     "PyYAML",
     "scipy",
     "numpy",
-    "slurmio"
+    "slurmio",
+    "tifffile",
+    "imio"
 ]
 
 license = {text = "MIT"}

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)
 

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)