Add tiled thresholds to 2d filtering for determining foreground pixels

benchmarks/filter_2d.py

@@ -50,6 +50,8 @@ def setup_filter(
         soma_diameter=settings.soma_diameter,
         log_sigma_size=settings.log_sigma_size,
         n_sds_above_mean_thresh=settings.n_sds_above_mean_thresh,
+        n_sds_above_mean_tiled_thresh=settings.n_sds_above_mean_tiled_thresh,
+        tiled_thresh_tile_size=settings.tiled_thresh_tile_size,
         torch_device=torch_device,
         dtype=settings.filtering_dtype.__name__,
         use_scipy=use_scipy,

cellfinder/core/detect/detect.py

@@ -53,6 +53,8 @@ def main(
     split_ball_z_size: float = 15,
     split_ball_overlap_fraction: float = 0.8,
     n_splitting_iter: int = 10,
+    n_sds_above_mean_tiled_thresh: float = 10,
+    tiled_thresh_tile_size: float | None = None,
     *,
     callback: Optional[Callable[[int], None]] = None,
 ) -> List[Cell]:
@@ -95,8 +97,8 @@ def main(
         Gaussian filter width (as a fraction of soma diameter) used during
         2d in-plane Laplacian of Gaussian filtering.
     n_sds_above_mean_thresh : float
-        Intensity threshold (the number of standard deviations above
-        the mean) of the filtered 2d planes used to mark pixels as
+        Per-plane intensity threshold (the number of standard deviations
+        above the mean) of the filtered 2d planes used to mark pixels as
         foreground or background.
     outlier_keep : bool, optional
         Whether to keep outliers during detection. Defaults to False.
@@ -129,6 +131,20 @@ def main(
         The number of iterations to run the 3d filtering on a cluster. Each
         iteration reduces the cluster size by the voxels not retained in
         the previous iteration.
+    n_sds_above_mean_tiled_thresh : float
+        Per-plane, per-tile intensity threshold (the number of standard
+        deviations above the mean) for the filtered 2d planes used to mark
+        pixels as foreground or background. When used, (tile size is not zero)
+        a pixel is marked as foreground if its intensity is above both the
+        per-plane and per-tile threshold. I.e. it's above the set number of
+        standard deviations of the per-plane average and of the per-plane
+        per-tile average for the tile that contains it.
+    tiled_thresh_tile_size : float
+        The tile size used to tile the x, y plane to calculate the local
+        average intensity for the tiled threshold. The value is multiplied
+        by soma diameter (i.e. 1 means one soma diameter). If zero or None, the
+        tiled threshold is disabled and only the per-plane threshold is used.
+        Tiling is done with 50% overlap when striding.
     callback : Callable[int], optional
         A callback function that is called every time a plane has finished
         being processed. Called with the plane number that has finished.
@@ -186,6 +202,8 @@ def main(
         ball_overlap_fraction=ball_overlap_fraction,
         log_sigma_size=log_sigma_size,
         n_sds_above_mean_thresh=n_sds_above_mean_thresh,
+        n_sds_above_mean_tiled_thresh=n_sds_above_mean_tiled_thresh,
+        tiled_thresh_tile_size=tiled_thresh_tile_size,
         outlier_keep=outlier_keep,
         artifact_keep=artifact_keep,
         save_planes=save_planes,
@@ -220,7 +238,9 @@ def main(
         plane_shape=settings.plane_shape,
         clipping_value=settings.clipping_value,
         threshold_value=settings.threshold_value,
-        n_sds_above_mean_thresh=n_sds_above_mean_thresh,
+        n_sds_above_mean_thresh=settings.n_sds_above_mean_thresh,
+        n_sds_above_mean_tiled_thresh=settings.n_sds_above_mean_tiled_thresh,
+        tiled_thresh_tile_size=settings.tiled_thresh_tile_size,
         log_sigma_size=log_sigma_size,
         soma_diameter=settings.soma_diameter,
         torch_device=torch_device,

cellfinder/core/detect/filters/plane/plane_filter.py

@@ -1,13 +1,12 @@
-from dataclasses import dataclass, field
 from typing import Tuple
 
 import torch
+import torch.nn.functional as F
 
 from cellfinder.core.detect.filters.plane.classical_filter import PeakEnhancer
 from cellfinder.core.detect.filters.plane.tile_walker import TileWalker
 
 
-@dataclass
 class TileProcessor:
     """
     Processor that filters each plane to highlight the peaks and also
@@ -63,19 +62,31 @@
     # voxels who are this many std above mean or more are set to
     # threshold_value
     n_sds_above_mean_thresh: float
+    # If used, voxels who are this many or more std above mean of the
+    # containing tile as well as above n_sds_above_mean_thresh for the plane
+    # average are set to threshold_value.
+    n_sds_above_mean_tiled_thresh: float
+    # the tile size, in pixels, that will be used to tile the x, y plane when
+    # we calculate the per-tile mean / std for use with
+    # n_sds_above_mean_tiled_thresh. We use 50% overlap when tiling.
+    local_threshold_tile_size_px: int = 0
+    # the torch device name
+    torch_device: str = ""
 
     # filter that finds the peaks in the planes
-    peak_enhancer: PeakEnhancer = field(init=False)
+    peak_enhancer: PeakEnhancer = None
     # generates tiles of the planes, with each tile marked as being inside
     # or outside the brain based on brightness
-    tile_walker: TileWalker = field(init=False)
+    tile_walker: TileWalker = None
 
     def __init__(
         self,
         plane_shape: Tuple[int, int],
         clipping_value: int,
         threshold_value: int,
         n_sds_above_mean_thresh: float,
+        n_sds_above_mean_tiled_thresh: float,
+        tiled_thresh_tile_size: float | None,
         log_sigma_size: float,
         soma_diameter: int,
         torch_device: str,
@@ -85,6 +96,12 @@
         self.clipping_value = clipping_value
         self.threshold_value = threshold_value
         self.n_sds_above_mean_thresh = n_sds_above_mean_thresh
+        self.n_sds_above_mean_tiled_thresh = n_sds_above_mean_tiled_thresh
+        if tiled_thresh_tile_size:
+            self.local_threshold_tile_size_px = int(
+                round(soma_diameter * tiled_thresh_tile_size)
+            )
+        self.torch_device = torch_device
 
         laplace_gaussian_sigma = log_sigma_size * soma_diameter
         self.peak_enhancer = PeakEnhancer(
@@ -131,7 +148,10 @@
             planes,
             enhanced_planes,
             self.n_sds_above_mean_thresh,
+            self.n_sds_above_mean_tiled_thresh,
+            self.local_threshold_tile_size_px,
             self.threshold_value,
+            self.torch_device,
         )
 
         return planes, inside_brain_tiles
@@ -145,21 +165,98 @@
     planes: torch.Tensor,
     enhanced_planes: torch.Tensor,
     n_sds_above_mean_thresh: float,
+    n_sds_above_mean_tiled_thresh: float,
+    local_threshold_tile_size_px: int,
     threshold_value: int,
+    torch_device: str,
 ) -> None:
     """
     Sets each plane (in-place) to threshold_value, where the corresponding
     enhanced_plane > mean + n_sds_above_mean_thresh*std. Each plane will be
     set to zero elsewhere.
     """
-    planes_1d = enhanced_planes.view(enhanced_planes.shape[0], -1)
+    z, y, x = enhanced_planes.shape
 
+    # ---- get per-plane global threshold ----
+    planes_1d = enhanced_planes.view(z, -1)
     # add back last dim
-    avg = torch.mean(planes_1d, dim=1, keepdim=True).unsqueeze(2)
-    sd = torch.std(planes_1d, dim=1, keepdim=True).unsqueeze(2)
-    threshold = avg + n_sds_above_mean_thresh * sd
+    std, mean = torch.std_mean(planes_1d, dim=1, keepdim=True)
+    threshold = mean.unsqueeze(2) + n_sds_above_mean_thresh * std.unsqueeze(2)
+    above_global = enhanced_planes > threshold
+
+    # ---- calculate the local tiled threshold ----
+    # we do 50% overlap so there's no jumps at boundaries
+    stride = local_threshold_tile_size_px // 2
+    # make tile even for ease of computation
+    tile_size = stride * 2
+    # Due to 50% overlap, to get tiles we move the tile by half tile (stride).
+    # Total moves will be y // stride - 2 (we start already with mask on first
+    # tile). So add back 1 for the first tile. Partial tiles are dropped
+    n_y_tiles = max(y // stride - 1, 1) if stride else 1
+    n_x_tiles = max(x // stride - 1, 1) if stride else 1
+    do_tile_y = n_y_tiles >= 2
+    do_tile_x = n_x_tiles >= 2
+    # we want at least one axis to have at least two tiles
+    if local_threshold_tile_size_px >= 2 and (do_tile_y or do_tile_x):
+        # num edge pixels dropped b/c moving by stride would move tile off edge
+        y_rem = y % stride
+        x_rem = x % stride
+        enhanced_planes_raw = enhanced_planes
+        if do_tile_y:
+            enhanced_planes = enhanced_planes[:, y_rem // 2 :, :]
+        if do_tile_x:
+            enhanced_planes = enhanced_planes[:, :, x_rem // 2 :]
+
+        # add empty channel dim after z "batch" dim -> zcyx
+        enhanced_planes = enhanced_planes.unsqueeze(1)
+        # unfold makes it 3 dim, z, M, L. L is number of tiles, M is tile area
+        unfolded = F.unfold(
+            enhanced_planes,
+            (tile_size if do_tile_y else y, tile_size if do_tile_x else x),
+            stride=stride,
+        )
+        # average the tile areas, for each tile
+        std, mean = torch.std_mean(unfolded, dim=1, keepdim=True)
+        threshold = mean + n_sds_above_mean_tiled_thresh * std
+
+        # reshape it back into Y by X tiles, instead of YX being one dim
+        threshold = threshold.reshape((z, n_y_tiles, n_x_tiles))
+
+        # we need total size of n_tiles * stride + stride + rem for the
+        # original size. So we add 2 strides and then chop off the excess above
+        # rem. We center it because of 50% overlap, the first tile is actually
+        # centered in between the first two strides
+        offsets = [(0, y), (0, x)]
+        for dim, do_tile, n_tiles, n, rem in [
+            (1, do_tile_y, n_y_tiles, y, y_rem),
+            (2, do_tile_x, n_x_tiles, x, x_rem),
+        ]:
+            if do_tile:
+                repeats = (
+                    torch.ones(n_tiles, dtype=torch.int, device=torch_device)
+                    * stride
+                )
+                # add total of 2 additional strides
+                repeats[0] = 2 * stride
+                repeats[-1] = 2 * stride
+                output_size = (n_tiles + 2) * stride
+
+                threshold = threshold.repeat_interleave(
+                    repeats, dim=dim, output_size=output_size
+                )
+                # drop the excess we gained from padding rem to whole stride
+                offset = (stride - rem) // 2
+                offsets[dim - 1] = offset, n + offset
+
+        # can't use slice(...) objects in jit code so use actual indices
+        (a, b), (c, d) = offsets
+        threshold = threshold[:, a:b, c:d]
+
+        above_local = enhanced_planes_raw > threshold
+        above = torch.logical_and(above_global, above_local)
+    else:
+        above = above_global
 
-    above = enhanced_planes > threshold
     planes[above] = threshold_value
     # subsequent steps only care about the values that are set to threshold or
     # above in planes. We set values in *planes* to threshold based on the

cellfinder/core/detect/filters/setup_filters.py

@@ -133,11 +133,31 @@ class DetectionSettings:
 
     n_sds_above_mean_thresh: float = 10
     """
-    Intensity threshold (the number of standard deviations above
-    the mean) of the filtered 2d planes used to mark pixels as
+    Per-plane intensity threshold (the number of standard deviations
+    above the mean) of the 2d filtered planes used to mark pixels as
     foreground or background.
     """
 
+    n_sds_above_mean_tiled_thresh: float = 10
+    """
+    Per-plane, per-tile intensity threshold (the number of standard deviations
+    above the mean) for the filtered 2d planes used to mark pixels as
+    foreground or background. When used, (tile size is not zero) a pixel is
+    marked as foreground if its intensity is above both the per-plane and
+    per-tile threshold. I.e. it's above the set number of standard deviations
+    of the per-plane average and of the per-plane per-tile average for the tile
+    that contains it.
+    """
+
+    tiled_thresh_tile_size: float | None = None
+    """
+    The tile size used to tile the x, y plane to calculate the local average
+    intensity for the tiled threshold. The value is multiplied by soma
+    diameter (i.e. 1 means one soma diameter). If zero or None, the tiled
+    threshold is disabled and only the per-plane threshold is used. Tiling is
+    done with 50% overlap when striding.
+    """
+
     outlier_keep: bool = False
     """Whether to keep outlier structures during detection."""
 

cellfinder/core/main.py

@@ -41,6 +41,8 @@ def main(
     split_ball_z_size: float = 15,
     split_ball_overlap_fraction: float = 0.8,
     n_splitting_iter: int = 10,
+    n_sds_above_mean_tiled_thresh: float = 10,
+    tiled_thresh_tile_size: float | None = None,
     *,
     detect_callback: Optional[Callable[[int], None]] = None,
     classify_callback: Optional[Callable[[int], None]] = None,
@@ -93,8 +95,8 @@ def main(
         Gaussian filter width (as a fraction of soma diameter) used during
         2d in-plane Laplacian of Gaussian filtering.
     n_sds_above_mean_thresh : float
-        Intensity threshold (the number of standard deviations above
-        the mean) of the filtered 2d planes used to mark pixels as
+        Per-plane intensity threshold (the number of standard deviations
+        above the mean) of the filtered 2d planes used to mark pixels as
         foreground or background.
     soma_spread_factor : float
         Cell spread factor for determining the largest cell volume before
@@ -148,6 +150,20 @@ def main(
         The number of iterations to run the 3d filtering on a cluster. Each
         iteration reduces the cluster size by the voxels not retained in
         the previous iteration.
+    n_sds_above_mean_tiled_thresh : float
+        Per-plane, per-tile intensity threshold (the number of standard
+        deviations above the mean) for the filtered 2d planes used to mark
+        pixels as foreground or background. When used, (tile size is not zero)
+        a pixel is marked as foreground if its intensity is above both the
+        per-plane and per-tile threshold. I.e. it's above the set number of
+        standard deviations of the per-plane average and of the per-plane
+        per-tile average for the tile that contains it.
+    tiled_thresh_tile_size : float
+        The tile size used to tile the x, y plane to calculate the local
+        average intensity for the tiled threshold. The value is multiplied
+        by soma diameter (i.e. 1 means one soma diameter). If zero or None, the
+        tiled threshold is disabled and only the per-plane threshold is used.
+        Tiling is done with 50% overlap when striding.
     detect_callback : Callable[int], optional
         Called every time a plane has finished being processed during the
         detection stage. Called with the plane number that has finished.
@@ -165,19 +181,21 @@ def main(
         logger.info("Detecting cell candidates")
 
         points = detect.main(
-            signal_array,
-            start_plane,
-            end_plane,
-            voxel_sizes,
-            soma_diameter,
-            max_cluster_size,
-            ball_xy_size,
-            ball_z_size,
-            ball_overlap_fraction,
-            soma_spread_factor,
-            n_free_cpus,
-            log_sigma_size,
-            n_sds_above_mean_thresh,
+            signal_array=signal_array,
+            start_plane=start_plane,
+            end_plane=end_plane,
+            voxel_sizes=voxel_sizes,
+            soma_diameter=soma_diameter,
+            max_cluster_size=max_cluster_size,
+            ball_xy_size=ball_xy_size,
+            ball_z_size=ball_z_size,
+            ball_overlap_fraction=ball_overlap_fraction,
+            soma_spread_factor=soma_spread_factor,
+            n_free_cpus=n_free_cpus,
+            log_sigma_size=log_sigma_size,
+            n_sds_above_mean_thresh=n_sds_above_mean_thresh,
+            n_sds_above_mean_tiled_thresh=n_sds_above_mean_tiled_thresh,
+            tiled_thresh_tile_size=tiled_thresh_tile_size,
             batch_size=detection_batch_size,
             torch_device=torch_device,
             callback=detect_callback,