cellfinder 1.3.3__py3-none-any.whl → 1.4.0a0__py3-none-any.whl

cellfinder/core/detect/filters/volume/ball_filter.py

@@ -1,29 +1,35 @@
+import math
 from functools import lru_cache
+from typing import Optional
 
 import numpy as np
-from numba import njit, objmode, prange
-from numba.core import types
-from numba.experimental import jitclass
+import torch
+import torch.nn.functional as F
 
 from cellfinder.core.tools.array_operations import bin_mean_3d
 from cellfinder.core.tools.geometry import make_sphere
 
-DEBUG = False
 
-uint32_3d_type = types.uint32[:, :, :]
-bool_3d_type = types.bool_[:, :, :]
-float_3d_type = types.float64[:, :, :]
+class InvalidVolume(ValueError):
+    """
+    Raised when the volume passed to BallFilter is too small or does not meet
+    requirements.
+    """
+
+    pass
 
 
 @lru_cache(maxsize=50)
 def get_kernel(ball_xy_size: int, ball_z_size: int) -> np.ndarray:
-    # Create a spherical kernel.
-    #
-    # This is done by:
-    # 1. Generating a binary sphere at a resolution *upscale_factor* larger
-    #    than desired.
-    # 2. Downscaling the binary sphere to get a 'fuzzy' sphere at the
-    #    original intended scale
+    """
+    Create a spherical kernel.
+
+    This is done by:
+    1. Generating a binary sphere at a resolution *upscale_factor* larger
+       than desired.
+    2. Downscaling the binary sphere to get a 'fuzzy' sphere at the
+       original intended scale
+    """
     upscale_factor: int = 7
     upscaled_kernel_shape = (
         upscale_factor * ball_xy_size,
@@ -42,11 +48,11 @@ def get_kernel(ball_xy_size: int, ball_z_size: int) -> np.ndarray:
         upscaled_ball_radius,
         upscaled_ball_centre_position,
     )
-    sphere_kernel = sphere_kernel.astype(np.float64)
+    sphere_kernel = sphere_kernel.astype(np.float32)
     kernel = bin_mean_3d(
         sphere_kernel,
-        bin_height=upscale_factor,
         bin_width=upscale_factor,
+        bin_height=upscale_factor,
         bin_depth=upscale_factor,
     )
 
@@ -59,359 +65,351 @@ def get_kernel(ball_xy_size: int, ball_z_size: int) -> np.ndarray:
     return kernel
 
 
-# volume indices/size is 64 bit for very large brains(!)
-spec = [
-    ("ball_xy_size", types.uint32),
-    ("ball_z_size", types.uint32),
-    ("tile_step_width", types.uint64),
-    ("tile_step_height", types.uint64),
-    ("THRESHOLD_VALUE", types.uint32),
-    ("SOMA_CENTRE_VALUE", types.uint32),
-    ("overlap_fraction", types.float64),
-    ("overlap_threshold", types.float64),
-    ("middle_z_idx", types.uint32),
-    ("_num_z_added", types.uint32),
-    ("kernel", float_3d_type),
-    ("volume", uint32_3d_type),
-    ("inside_brain_tiles", bool_3d_type),
-]
-
-
-@jitclass(spec=spec)
 class BallFilter:
     """
     A 3D ball filter.
 
-    This runs a spherical kernel across the (x, y) dimensions
+    This runs a spherical kernel across the 2d planar dimensions
     of a *ball_z_size* stack of planes, and marks pixels in the middle
-    plane of the stack that have a high enough intensity within the
-    spherical kernel.
+    plane of the stack that have a high enough intensity over the
+    the spherical kernel.
+
+    Parameters
+    ----------
+    plane_height, plane_width : int
+        Height/width of the planes.
+    ball_xy_size : int
+        Diameter of the spherical kernel in the x/y dimensions.
+    ball_z_size : int
+        Diameter of the spherical kernel in the z dimension.
+        Equal to the number of planes stacked to filter
+        the central plane of the stack.
+    overlap_fraction : float
+        The fraction of pixels within the spherical kernel that
+        have to be over *threshold_value* for a pixel to be marked
+        as having a high intensity.
+    threshold_value : int
+        Value above which an individual pixel is considered to have
+        a high intensity.
+    soma_centre_value : int
+        Value used to mark pixels with a high enough intensity.
+    tile_height, tile_width : int
+        Width/height of individual tiles in the mask generated by
+        2D filtering.
+    dtype : str
+        The data-type of the input planes and the type to use internally.
+        E.g. "float32".
+    batch_size: int
+        The number of planes that will be typically passed in a single batch to
+        `append`. This is only used to calculate `num_batches_before_ready`.
+        Defaults to 1.
+    torch_device: str
+        The device on which the data and processing occurs on. Can be e.g.
+        "cpu", "cuda" etc. Defaults to "cpu". Any data passed to the filter
+        must be on this device. Returned data will also be on this device.
+    use_mask : bool
+        Whether tiling masks will be used in `append`. If False, tile masks
+        won't be passed in and/or will be ignored. Defaults to True.
+    """
+
+    num_batches_before_ready: int
     """
+    The number of batches of size `batch_size` passed to `append`
+    before `ready` would return True.
+    """
+
+    # the inside brain tiled mask, if tiles are used (use_mask is True)
+    inside_brain_tiles: Optional[torch.Tensor] = None
 
     def __init__(
         self,
-        plane_width: int,
         plane_height: int,
+        plane_width: int,
         ball_xy_size: int,
         ball_z_size: int,
         overlap_fraction: float,
-        tile_step_width: int,
-        tile_step_height: int,
         threshold_value: int,
         soma_centre_value: int,
+        tile_height: int,
+        tile_width: int,
+        dtype: str,
+        batch_size: int = 1,
+        torch_device: str = "cpu",
+        use_mask: bool = True,
     ):
-        """
-        Parameters
-        ----------
-        plane_width, plane_height :
-            Width/height of the planes.
-        ball_xy_size :
-            Diameter of the spherical kernel in the x/y dimensions.
-        ball_z_size :
-            Diameter of the spherical kernel in the z dimension.
-            Equal to the number of planes that stacked to filter
-            the central plane of the stack.
-        overlap_fraction :
-            The fraction of pixels within the spherical kernel that
-            have to be over *threshold_value* for a pixel to be marked
-            as having a high intensity.
-        tile_step_width, tile_step_height :
-            Width/height of individual tiles in the mask generated by
-            2D filtering.
-        threshold_value :
-            Value above which an individual pixel is considered to have
-            a high intensity.
-        soma_centre_value :
-            Value used to mark pixels with a high enough intensity.
-        """
         self.ball_xy_size = ball_xy_size
         self.ball_z_size = ball_z_size
         self.overlap_fraction = overlap_fraction
-        self.tile_step_width = tile_step_width
-        self.tile_step_height = tile_step_height
+        self.tile_step_height = tile_height
+        self.tile_step_width = tile_width
+
+        d1 = plane_height
+        d2 = plane_width
+        ball_xy_size = self.ball_xy_size
+        if d1 < ball_xy_size or d2 < ball_xy_size:
+            raise InvalidVolume(
+                f"Invalid plane size {d1}x{d2}. Needs to be at least "
+                f"{ball_xy_size} in each dimension"
+            )
 
         self.THRESHOLD_VALUE = threshold_value
         self.SOMA_CENTRE_VALUE = soma_centre_value
 
-        # getting kernel is not jitted
-        with objmode(kernel=float_3d_type):
-            kernel = get_kernel(ball_xy_size, ball_z_size)
-        self.kernel = kernel
-
-        self.overlap_threshold = np.sum(self.overlap_fraction * self.kernel)
+        # kernel comes in as XYZ, change to ZYX so it aligns with data
+        kernel = np.moveaxis(get_kernel(ball_xy_size, self.ball_z_size), 2, 0)
+        self.overlap_threshold = np.sum(self.overlap_fraction * kernel)
+        self.kernel_xy_size = kernel.shape[-2:]
+        self.kernel_z_size = self.ball_z_size
+
+        # convert to right type and pin for faster copying
+        kernel = torch.from_numpy(kernel).type(getattr(torch, dtype))
+        if torch_device != "cpu":
+            # torch at one point threw a cuda memory error when splitting cells
+            # on cpu due to pinning. It's best to only pin on using cuda
+            kernel.pin_memory()
+        # add 2 dimensions at the start so we have 11ZYX We need this shape in
+        # the conv step
+        self.kernel = (
+            kernel.unsqueeze(0)
+            .unsqueeze(0)
+            .to(device=torch_device, non_blocking=True)
+        )
 
-        # Stores the current planes that are being filtered
-        # first axis is z for faster rotating the z-axis
-        self.volume = np.empty(
-            (ball_z_size, plane_width, plane_height),
-            dtype=np.uint32,
+        self.num_batches_before_ready = int(
+            math.ceil(self.ball_z_size / batch_size)
+        )
+        # Stores the current planes that are being filtered. Start with no data
+        self.volume = torch.empty(
+            (0, plane_height, plane_width),
+            dtype=getattr(torch, dtype),
         )
         # Index of the middle plane in the volume
-        self.middle_z_idx = int(np.floor(ball_z_size / 2))
-        self._num_z_added = 0
+        self.middle_z_idx = int(np.floor(self.ball_z_size / 2))
 
+        if not use_mask:
+            return
         # first axis is z
-        self.inside_brain_tiles = np.empty(
+        n_vertical_tiles = int(np.ceil(plane_height / self.tile_step_height))
+        n_horizontal_tiles = int(np.ceil(plane_width / self.tile_step_width))
+        # Stores tile masks. We start with no data
+        self.inside_brain_tiles = torch.empty(
             (
-                ball_z_size,
-                int(np.ceil(plane_width / tile_step_width)),
-                int(np.ceil(plane_height / tile_step_height)),
+                0,
+                n_vertical_tiles,
+                n_horizontal_tiles,
             ),
-            dtype=np.bool_,
+            dtype=torch.bool,
         )
 
     @property
-    def ready(self) -> bool:
+    def first_valid_plane(self) -> int:
         """
-        Return `True` if enough planes have been appended to run the filter.
+        The index in `self.volume` (or the planes passed in) that will be the
+        first plane returned from `get_processed_planes`.
+
+        E.g. if `ball_z_size` is 3, then this may return 1. Meaning the second
+        plane passed to `append` (index 1), will be the first returned plane
+        by `get_processed_planes`.
         """
-        return self._num_z_added >= self.ball_z_size
+        return int(math.floor(self.ball_z_size / 2))
 
-    def append(self, plane: np.ndarray, mask: np.ndarray) -> None:
+    @property
+    def remaining_planes(self) -> int:
         """
-        Add a new 2D plane to the filter.
+        The number of planes in `self.volume` (or the planes passed in) that
+        will remain unprocessed after all the planes have been `walk`ed
+        and `get_processed_planes` called.
+
+        E.g. if `ball_z_size` is 3, then this may return 1. Meaning the last
+        plane passed to `append`, will never be returned by
+        `get_processed_planes` because the filter "center" never overlapped
+        with it.
         """
-        if DEBUG:
-            assert [e for e in plane.shape[:2]] == [
-                e for e in self.volume.shape[1:]
-            ], 'plane shape mismatch, expected "{}", got "{}"'.format(
-                [e for e in self.volume.shape[1:]],
-                [e for e in plane.shape[:2]],
-            )
-            assert [e for e in mask.shape[:2]] == [
-                e for e in self.inside_brain_tiles.shape[1:]
-            ], 'mask shape mismatch, expected"{}", got {}"'.format(
-                [e for e in self.inside_brain_tiles.shape[1:]],
-                [e for e in mask.shape[:2]],
-            )
+        return self.ball_z_size - self.first_valid_plane - 1
 
-        if self.ready:
-            # Shift everything down by one to make way for the new plane
-            # this is faster than np.roll, especially with z-axis first
-            self.volume[:-1, :, :] = self.volume[1:, :, :]
-            self.inside_brain_tiles[:-1, :, :] = self.inside_brain_tiles[
-                1:, :, :
-            ]
+    @property
+    def ready(self) -> bool:
+        """
+        Return whether enough planes have been appended to run the filter
+        using `walk`.
+        """
+        return self.volume.shape[0] >= self.kernel_z_size
+
+    def append(
+        self, planes: torch.Tensor, masks: Optional[torch.Tensor] = None
+    ) -> None:
+        """
+        Add a new z-stack to the filter.
 
-        # index for *next* slice is num we added *so far* until max
-        idx = min(self._num_z_added, self.ball_z_size - 1)
-        self._num_z_added += 1
+        Previous stacks passed to `append` are removed, except enough planes
+        at the top of the previous z-stack to provide padding so we can filter
+        starting from the first plane in `planes`. The first time we call
+        `append`, `first_valid_plane` is the first plane to actually be
+        filtered in the z-stack due to lack of padding.
 
-        # Add the new plane to the top of volume and inside_brain_tiles
-        self.volume[idx, :, :] = plane
-        self.inside_brain_tiles[idx, :, :] = mask
+        So make sure to call `walk`/`get_processed_planes` before calling
+        `append` again.
 
-    def get_middle_plane(self) -> np.ndarray:
-        """
-        Get the plane in the middle of self.volume.
+        Parameters
+        ----------
+        planes : torch.Tensor
+            The z-stack. There can be one or more planes in the stack, but it
+            must have 3 dimensions. Input data is not modified.
+        masks : torch.Tensor
+            A z-stack tile mask, indicating for each tile whether it's in or
+            outside the brain. If the latter it's excluded.
+
+            If `use_mask` was True, this must be provided. If False, this
+            parameter will be ignored.
+
+            Input data is not modified.
         """
-        return self.volume[self.middle_z_idx, :, :].copy()
-
-    def walk(self, parallel: bool = False) -> None:
-        # **don't** pass parallel as keyword arg - numba struggles with it
-        # Highly optimised because most time critical
-        ball_radius = self.ball_xy_size // 2
-        # Get extents of image that are covered by tiles
-        tile_mask_covered_img_width = (
-            self.inside_brain_tiles.shape[1] * self.tile_step_width
-        )
-        tile_mask_covered_img_height = (
-            self.inside_brain_tiles.shape[2] * self.tile_step_height
-        )
-        # Get maximum offsets for the ball
-        max_width = tile_mask_covered_img_width - self.ball_xy_size
-        max_height = tile_mask_covered_img_height - self.ball_xy_size
-
-        # we have to pass the raw volume so walk doesn't use its edits as it
-        # processes the volume. self.volume is the one edited in place
-        input_volume = self.volume.copy()
-
-        if parallel:
-            _walk_parallel(
-                max_height,
-                max_width,
-                self.tile_step_width,
-                self.tile_step_height,
-                self.inside_brain_tiles,
-                input_volume,
-                self.volume,
-                self.kernel,
-                ball_radius,
-                self.middle_z_idx,
-                self.overlap_threshold,
-                self.THRESHOLD_VALUE,
-                self.SOMA_CENTRE_VALUE,
+        if self.volume.shape[0]:
+            if self.volume.shape[0] < self.kernel_z_size:
+                num_remaining_with_padding = 0
+            else:
+                num_remaining = self.kernel_z_size - (self.middle_z_idx + 1)
+                num_remaining_with_padding = num_remaining + self.middle_z_idx
+
+            self.volume = torch.cat(
+                [self.volume[-num_remaining_with_padding:, :, :], planes],
+                dim=0,
             )
+
+            if self.inside_brain_tiles is not None:
+                self.inside_brain_tiles = torch.cat(
+                    [
+                        self.inside_brain_tiles[
+                            -num_remaining_with_padding:, :, :
+                        ],
+                        masks,
+                    ],
+                    dim=0,
+                )
         else:
-            _walk_single(
-                max_height,
-                max_width,
-                self.tile_step_width,
-                self.tile_step_height,
-                self.inside_brain_tiles,
-                input_volume,
-                self.volume,
-                self.kernel,
-                ball_radius,
-                self.middle_z_idx,
-                self.overlap_threshold,
-                self.THRESHOLD_VALUE,
-                self.SOMA_CENTRE_VALUE,
-            )
+            self.volume = planes.clone()
+            if self.inside_brain_tiles is not None:
+                self.inside_brain_tiles = masks.clone()
 
+    def get_processed_planes(self) -> np.ndarray:
+        """
+        After passing enough planes to `append`, and after `walk`, this returns
+        a copy of the processed planes as a numpy z-stack.
 
-@njit(cache=True)
-def _cube_overlaps(
-    volume: np.ndarray,
-    x_start: int,
-    x_end: int,
-    y_start: int,
-    y_end: int,
-    overlap_threshold: float,
-    threshold_value: int,
-    kernel: np.ndarray,
-) -> bool:  # Highly optimised because most time critical
-    """
-    For each pixel in cube in volume that is greater than THRESHOLD_VALUE, sum
-    up the corresponding pixels in *kernel*. If the total is less than
-    overlap_threshold, return False, otherwise return True.
+        It only starts returning planes corresponding to plane
+        `first_valid_plane` relative to the first planes passed to `append`.
+        E.g. if `ball_z_size` is 3 and `first_valid_plane` is 1, and you passed
+        5 planes total to `append`, then this will have returned planes [1, 3].
 
-    Halfway through scanning the z-planes, if the total overlap is
-    less than 0.4 * overlap_threshold, this will return False early
-    without scanning the second half of the z-planes.
+        Notice the last plane was not included, because we return only "middle"
+        planes - planes that can correspond to the center of a ball.
+        """
+        if not self.ready:
+            raise TypeError("Not enough planes were appended")
+
+        num_processed = self.volume.shape[0] - self.kernel_z_size + 1
+        assert num_processed
+        middle = self.middle_z_idx
+        planes = (
+            self.volume[middle : middle + num_processed, :, :]
+            .cpu()
+            .numpy()
+            .copy()
+        )
+        return planes
 
-    Parameters
-    ----------
-    volume :
-        3D array.
-    x_start, x_end, y_start, y_end :
-        The start and end indices in volume that form the cube. End is
-        exclusive
-    overlap_threshold :
-        Threshold above which to return True.
-    threshold_value :
-        Value above which a pixel is marked as being part of a cell.
-    kernel :
-        3D array, with the same shape as *cube* in the volume.
-    """
-    current_overlap_value = 0.0
-
-    middle = np.floor(volume.shape[0] / 2) + 1
-    halfway_overlap_thresh = (
-        overlap_threshold * 0.4
-    )  # FIXME: do not hard code value
-
-    for z in range(volume.shape[0]):
-        # TODO: OPTIMISE: step from middle to outer boundaries to check
-        # more data first
-        #
-        # If halfway through the array, and the overlap value isn't more than
-        # 0.4 * the overlap threshold, return
-        if z == middle and current_overlap_value < halfway_overlap_thresh:
-            return False  # DEBUG: optimisation attempt
-
-        for y in range(y_start, y_end):
-            for x in range(x_start, x_end):
-                # includes self.SOMA_CENTRE_VALUE
-                if volume[z, x, y] >= threshold_value:
-                    # x/y must be shifted in kernel because we x/y is relative
-                    # to the full volume, so shift it to relative to the cube
-                    current_overlap_value += kernel[
-                        x - x_start, y - y_start, z
-                    ]
-
-    return current_overlap_value > overlap_threshold
-
-
-@njit
-def _is_tile_to_check(
-    x: int,
-    y: int,
-    middle_z: int,
-    tile_step_width: int,
-    tile_step_height: int,
-    inside_brain_tiles: np.ndarray,
-) -> bool:  # Highly optimised because most time critical
-    """
-    Check if the tile containing pixel (x, y) is a tile that needs checking.
-    """
-    x_in_mask = x // tile_step_width  # TEST: test bounds (-1 range)
-    y_in_mask = y // tile_step_height  # TEST: test bounds (-1 range)
-    return inside_brain_tiles[middle_z, x_in_mask, y_in_mask]
+    def walk(self) -> None:
+        """
+        Applies the filter to all the planes passed to `append`.
 
+        May only be called if `ready` was True.
 
-def _walk_base(
-    max_height: int,
-    max_width: int,
-    tile_step_width: int,
+        You can get the processed planes from `get_processed_planes`.
+        """
+        if not self.ready:
+            raise TypeError("Called walk before enough planes were appended")
+
+        _walk(
+            self.kernel_z_size,
+            self.middle_z_idx,
+            self.tile_step_height,
+            self.tile_step_width,
+            self.overlap_threshold,
+            self.THRESHOLD_VALUE,
+            self.SOMA_CENTRE_VALUE,
+            self.kernel,
+            self.volume,
+            self.inside_brain_tiles,
+        )
+
+
+@torch.jit.script
+def _walk(
+    kernel_z_size: int,
+    middle: int,
     tile_step_height: int,
-    inside_brain_tiles: np.ndarray,
-    input_volume: np.ndarray,
-    volume: np.ndarray,
-    kernel: np.ndarray,
-    ball_radius: int,
-    middle_z: int,
+    tile_step_width: int,
     overlap_threshold: float,
     threshold_value: int,
     soma_centre_value: int,
-) -> None:
-    """
-    Scan through *volume*, and mark pixels where there are enough surrounding
-    pixels with high enough intensity.
+    kernel: torch.Tensor,
+    volume: torch.Tensor,
+    inside_brain_tiles: Optional[torch.Tensor],
+):
+    num_process = volume.shape[0] - kernel_z_size + 1
+    height, width = volume.shape[1:]
+    num_z = kernel.shape[2]
+
+    # threshold volume so it's zero/one. And add two dims at start so
+    # it's 11ZYX
+    volume_tresh = (
+        (volume >= threshold_value)
+        .unsqueeze(0)
+        .unsqueeze(0)
+        .type(kernel.dtype)
+    )
 
-    The surrounding area is defined by the *kernel*.
+    # we do a plane at a time, volume: i:i+num_z, for plane i+middle
+    for i in range(num_process):
+        # spherical kernel is symmetric so convolution=correlation. Use
+        # binary threshold mask over the kernel to sum the value of the
+        # kernel at voxels that are bright
+        overlap = F.conv3d(
+            volume_tresh[:, :, i : i + num_z, :, :],
+            kernel,
+            stride=1,
+            padding="valid",
+        )[0, 0, 0, :, :]
+        overlaps = overlap > overlap_threshold
+
+        # only edit the volume that is valid - conv excludes edges so we
+        # only edit the plane parts returned by conv3d
+        height_valid, width_valid = overlaps.shape
+        height_offset = (height - height_valid) // 2
+        width_offset = (width - width_valid) // 2
+        sub_volume = volume[
+            i + middle,
+            height_offset : height_offset + height_valid,
+            width_offset : width_offset + width_valid,
+        ]
+
+        # do we use tile masks or just overlapping?
+        if inside_brain_tiles is not None:
+            # unfold tiles to cover the full voxel area each tile covers
+            inside = (
+                inside_brain_tiles[i + middle, :, :]
+                .repeat_interleave(tile_step_height, dim=0)
+                .repeat_interleave(tile_step_width, dim=1)
+            )
+            # again only process pixels in the valid area
+            inside = inside[
+                height_offset : height_offset + height_valid,
+                width_offset : width_offset + width_valid,
+            ]
 
-    Parameters
-    ----------
-    max_height, max_width :
-        Maximum offsets for the ball filter.
-    inside_brain_tiles :
-        3d array containing information on whether a tile is
-        inside the brain or not. Tiles outside the brain are skipped.
-    input_volume :
-        3D array containing the plane-filtered data passed to the function
-        before walking. volume is edited in place, so this is the original
-        volume to prevent the changes for some cubes affective other cubes
-        during a single walk call.
-    volume :
-        3D array containing the plane-filtered data - edited in place.
-    kernel :
-        3D array
-    ball_radius :
-        Radius of the ball in the xy plane.
-    soma_centre_value :
-        Value that is used to mark pixels in *volume*.
-
-    Notes
-    -----
-    Warning: modifies volume in place!
-    """
-    for y in prange(max_height):
-        for x in prange(max_width):
-            ball_centre_x = x + ball_radius
-            ball_centre_y = y + ball_radius
-            if _is_tile_to_check(
-                ball_centre_x,
-                ball_centre_y,
-                middle_z,
-                tile_step_width,
-                tile_step_height,
-                inside_brain_tiles,
-            ):
-                if _cube_overlaps(
-                    input_volume,
-                    x,
-                    x + kernel.shape[0],
-                    y,
-                    y + kernel.shape[1],
-                    overlap_threshold,
-                    threshold_value,
-                    kernel,
-                ):
-                    volume[middle_z, ball_centre_x, ball_centre_y] = (
-                        soma_centre_value
-                    )
-
-
-_walk_parallel = njit(parallel=True)(_walk_base)
-_walk_single = njit(parallel=False)(_walk_base)
+            # must have enough ball overlap to be bright/tile is in brain
+            sub_volume[torch.logical_and(overlaps, inside)] = soma_centre_value
+
+        else:
+            # must have enough ball overlap to be bright
+            sub_volume[overlaps] = soma_centre_value