Fix ROI masks and other things

djimaging/autorois/autoshift_utils.py

@@ -139,15 +139,18 @@ def compute_corr_map_match_indexes(
     # Iterate over all possible shifts in both x and y directions within the range [-ShiftMax, ShiftMax]
     for i, (dx, dy) in enumerate(dx_dy_pairs, start=1):
         # Crop the shifted_video and stack_corr images to the same size
-        cropped = corr_map[max(0, -dx):min(ref_corr_map.shape[0], ref_corr_map.shape[0] - dx),
-                  max(0, -dy):min(ref_corr_map.shape[1], ref_corr_map.shape[1] - dy)]
+        cropped = corr_map[max(0, -dx):min(corr_map.shape[0], corr_map.shape[0] - dx),
+                  max(0, -dy):min(corr_map.shape[1], corr_map.shape[1] - dy)]
 
         ref_cropped = ref_corr_map[max(0, dx):min(ref_corr_map.shape[0], ref_corr_map.shape[0] + dx),
                       max(0, dy):min(ref_corr_map.shape[1], ref_corr_map.shape[1] + dy)]
 
         # Calculate the correlation coefficient between the cropped images
         if metric == 'corr':
-            score = np.corrcoef(cropped.flatten(), ref_cropped.flatten())[0, 1]
+            a = cropped.ravel() - cropped.mean()
+            b = ref_cropped.ravel() - ref_cropped.mean()
+            denom = np.linalg.norm(a) * np.linalg.norm(b)
+            score = np.dot(a, b) / denom if denom > 0 else 0.0
         elif metric == 'mse':
             score = -np.mean((cropped - ref_cropped) ** 2)
         else:
@@ -165,6 +168,7 @@ def compute_corr_map_match_indexes(
             ax2.imshow(cropped)
 
             plt.show()
+            plt.close(fig)
 
         # Store the correlation coefficient in the cross-correlation matrix
         match_indexes[dx + shift_max, dy + shift_max] = score
@@ -176,9 +180,9 @@ def compute_corr_map_match_indexes(
     v_max = np.max(match_indexes)
     v_max_row_loc, v_max_col_loc = np.unravel_index(np.argmax(match_indexes), match_indexes.shape)
 
-    # Shift the ROI mask
-    x_shift = v_max_row_loc - shift_max
-    y_shift = v_max_col_loc - shift_max
+    # Recover the shift: dx stored at index dx+shift_max, so dx = index-shift_max
+    x_shift = shift_max - v_max_row_loc
+    y_shift = shift_max - v_max_col_loc
 
     if verbose:
         print("ROI mask shifted by", x_shift, "pixels in x and", y_shift, "pixels in y")

djimaging/autorois/cellpose_wrapper.py

@@ -30,7 +30,7 @@ def __init__(self, init_kwargs: dict, eval_kwargs: dict) -> None:
 
         from cellpose import models, io
         io.logger_setup()
-        self.model = models.Cellpose(**self.init_kwargs)
+        self.model = models.CellposeModel(**self.init_kwargs)
 
     @staticmethod
     def stack_to_image(stack: np.ndarray, n_artifact: int) -> np.ndarray:
@@ -139,7 +139,7 @@ def create_mask_from_data(
         if multiple_stacks and 'stitch_threshold' not in eval_kwargs:
             eval_kwargs['stitch_threshold'] = 1e-9  # Stitch everything together
 
-        masks, flows, styles, diams = self.model.eval(imgs, do_3D=do_3D, **eval_kwargs)
+        masks, flows, *_ = self.model.eval(imgs, do_3D=do_3D, **eval_kwargs)
 
         if plot:
             self.plot_results(imgs, masks, flows)

djimaging/autorois/corr_roi_mask_utils.py

@@ -1,6 +1,5 @@
+import heapq
 import warnings
-from collections import deque
-from typing import Callable
 
 import numpy as np
 from matplotlib import pyplot as plt
@@ -50,8 +49,9 @@ class CorrRoiMask:
         If ``True``, use channel 0; otherwise use channel 1. Default is
         ``True``.
     grow_use_corr_map : bool, optional
-        If ``True``, use the precomputed correlation map during growing instead
-        of computing pairwise correlations on the fly. Default is ``False``.
+        If ``True``, precompute a mean-centered stack and use vectorised
+        seed-to-pixel correlation during growing instead of computing
+        pairwise correlations on the fly. Default is ``False``.
     grow_threshold : float or None, optional
         Correlation threshold used during growing. If ``None``, the
         line-threshold is used. Default is ``None``.
@@ -119,22 +119,28 @@ def create_mask_from_data(
         if len(kwargs) > 0:
             warnings.warn(f'ignoring kwargs: {kwargs}')
 
-        cut_x = np.array(self.cut_x)
-        cut_x[0] = np.maximum(cut_x[0], n_artifact)
+        cut_x = (max(self.cut_x[0], n_artifact), self.cut_x[1])
 
         stack = ch0_stack if self.use_ch0_stack else ch1_stack
 
-        n_pix_min = self.n_pix_min or 2
-        n_pix_min = int(np.maximum(np.round(self.min_area_um2 / (pixel_size_um[0] * pixel_size_um[1])), n_pix_min))
+        if stack is None:
+            raise ValueError("ch1_stack must be provided when use_ch0_stack=False")
 
-        n_pix_max = self.n_pix_max or 2
-        n_pix_max = int(np.maximum(np.round(self.max_area_um2 / (pixel_size_um[0] * pixel_size_um[1])), n_pix_max))
+        if self.n_pix_min is not None:
+            n_pix_min = self.n_pix_min
+        else:
+            n_pix_min = max(2, int(np.round(self.min_area_um2 / (pixel_size_um[0] * pixel_size_um[1]))))
 
-        # At most twice as long as wide
-        n_pix_max_x = int(
-            np.maximum(np.sqrt(2 * self.max_area_um2), 1)) if self.n_pix_max_x is None else self.n_pix_max_x
-        n_pix_max_z = int(
-            np.maximum(np.sqrt(2 * self.max_area_um2), 1)) if self.n_pix_max_z is None else self.n_pix_max_z
+        if self.n_pix_max is not None:
+            n_pix_max = self.n_pix_max
+        else:
+            n_pix_max = max(2, int(np.round(self.max_area_um2 / (pixel_size_um[0] * pixel_size_um[1]))))
+
+        # At most twice as long as wide, converted to pixels
+        n_pix_max_x = max(1, int(np.sqrt(2 * self.max_area_um2) / pixel_size_um[
+            0])) if self.n_pix_max_x is None else self.n_pix_max_x
+        n_pix_max_z = max(1, int(np.sqrt(2 * self.max_area_um2) / pixel_size_um[
+            1])) if self.n_pix_max_z is None else self.n_pix_max_z
 
         print(f'Creating ROI mask for stack of shape={stack.shape}, '
               f'with n_pix_min={n_pix_min}, n_pix_max={n_pix_max}, '
@@ -233,7 +239,7 @@ def corr_map_rois(
         n_pix_max_x: int = 5,
         n_pix_max_z: int = 5,
         n_pix_max: int = 25,
-        n_pix_min: float = 3.,
+        n_pix_min: int = 3,
         only_local_thresh_pixels: bool = False,
         corr_map: np.ndarray | None = None,
         grow_use_corr_map: bool = False,
@@ -264,7 +270,7 @@ def corr_map_rois(
         Maximum ROI extent in z (pixels). Default is 5.
     n_pix_max : int, optional
         Maximum total number of pixels per ROI. Default is 25.
-    n_pix_min : float, optional
+    n_pix_min : int, optional
         Minimum number of pixels for a ROI to be kept. Default is 3.
     only_local_thresh_pixels : bool, optional
         If ``True``, only pixels above the threshold can be added during
@@ -273,8 +279,10 @@ def corr_map_rois(
         Pre-computed correlation map. Computed from ``stack`` if ``None``.
         Default is ``None``.
     grow_use_corr_map : bool, optional
-        If ``True``, use ``corr_map`` values during growing. Default is
-        ``False``.
+        If ``True``, precompute a mean-centered stack and per-pixel norms
+        once, then use a vectorised seed-to-all-pixel correlation map during
+        growing instead of computing per-neighbour correlations on the fly.
+        Default is ``False``.
     line_threshold : np.ndarray or None, optional
         Pre-computed per-z-line threshold. Computed if ``None``. Default is
         ``None``.
@@ -301,11 +309,11 @@ def corr_map_rois(
             corr_map, cut_x=cut_x, cut_z=cut_z, q=line_threshold_q, line_threshold_min=line_threshold_min)
 
     nx, nz, nt = stack.shape
-    seed_ixs, seed_iys = np.argwhere(corr_map >= line_threshold).T
+    seed_ixs, seed_izs = np.argwhere(corr_map >= line_threshold).T
     # Sort by neighborhood correlation
-    sort_idxs = np.argsort(corr_map[seed_ixs, seed_iys])[::-1]
+    sort_idxs = np.argsort(corr_map[seed_ixs, seed_izs])[::-1]
     seed_ixs = seed_ixs[sort_idxs]
-    seed_iys = seed_iys[sort_idxs]
+    seed_izs = seed_izs[sort_idxs]
 
     roi_mask = np.zeros((nx, nz), dtype=int)
 
@@ -320,23 +328,32 @@ def corr_map_rois(
 
     p_threshold = grow_threshold if grow_threshold is not None else line_threshold
 
-    for seed_ix, seed_iz in zip(seed_ixs, seed_iys):
+    # Precompute centered stack and norms for efficient seed-to-pixel correlation
+    if grow_use_corr_map:
+        stack_float = stack.astype(float)
+        stack_centered = stack_float - stack_float.mean(axis=2, keepdims=True)
+        stack_norms = np.sqrt((stack_centered ** 2).sum(axis=2))
+    else:
+        stack_centered = None
+        stack_norms = None
+
+    for seed_ix, seed_iz in zip(seed_ixs, seed_izs):
         if roi_mask[seed_ix, seed_iz] != 0:
             continue
 
         roi_pixels = corr_map_grow_roi(
             seed_ix, seed_iz, n_pix_max=n_pix_max, n_pix_max_x=n_pix_max_x, n_pix_max_z=n_pix_max_z,
             p_threshold=p_threshold, stack=stack, ignore_map=ignore_map,
-            corr_map=corr_map if grow_use_corr_map else None)
+            stack_centered=stack_centered, stack_norms=stack_norms)
 
         if len(roi_pixels) == 0:
             continue
 
-        roi_ixs, roi_iys = np.array(roi_pixels).T
+        roi_ixs, roi_izs = np.array(roi_pixels).T
 
         if roi_ixs.size >= n_pix_min:
-            roi_mask[roi_ixs, roi_iys] = roi_id
-            ignore_map[roi_ixs, roi_iys] = True
+            roi_mask[roi_ixs, roi_izs] = roi_id
+            ignore_map[roi_ixs, roi_izs] = True
             roi_id += 1
 
     if plot:
@@ -398,7 +415,8 @@ def corr_map_grow_roi(
         p_threshold: np.ndarray | float,
         stack: np.ndarray,
         ignore_map: np.ndarray | None = None,
-        corr_map: np.ndarray | None = None,
+        stack_centered: np.ndarray | None = None,
+        stack_norms: np.ndarray | None = None,
 ) -> list:
     """Grow a single ROI from a seed pixel by collecting correlated neighbours.
 
@@ -423,14 +441,20 @@ def corr_map_grow_roi(
         Pixels below the threshold are not added to the ROI.
     stack : np.ndarray
         3-D imaging stack of shape (nx, nz, nt) used for pairwise correlation
-        when ``corr_map`` is ``None``.
+        when ``stack_centered`` is ``None``.
     ignore_map : np.ndarray or None, optional
         Boolean 2-D mask of shape (nx, nz); ``True`` pixels are skipped.
         Default is ``None``.
-    corr_map : np.ndarray or None, optional
-        Pre-computed 2-D correlation map. If provided, its values are used
-        instead of computing pairwise correlations on the fly. Default is
-        ``None``.
+    stack_centered : np.ndarray or None, optional
+        Pre-computed mean-centered stack of shape (nx, nz, nt).  When provided
+        together with ``stack_norms``, a vectorised seed-to-all-pixel
+        correlation map is computed at the start of the call instead of
+        evaluating per-neighbour correlations on the fly.
+        Default is ``None``.
+    stack_norms : np.ndarray or None, optional
+        Pre-computed per-pixel L2 norms of the centered traces, shape
+        (nx, nz).  Required when ``stack_centered`` is given.
+        Default is ``None``.
 
     Returns
     -------
@@ -439,24 +463,52 @@ def corr_map_grow_roi(
     """
     nx, nz, nt = stack.shape
     visited = np.zeros((nx, nz), dtype=bool)
+    in_heap = np.zeros((nx, nz), dtype=bool)
     if ignore_map is not None:
         visited[ignore_map] = True
+        in_heap[ignore_map] = True
 
     p_threshold = np.asarray(p_threshold)
     if p_threshold.size == 1:
         p_threshold = np.full(nz, p_threshold)
 
-    queue = deque([(seed_ix, seed_iz, 1.)])
+    # Compute seed-to-all-pixel correlation map (vectorised) or prepare
+    # for on-the-fly per-neighbour computation.
+    seed_centered = None
+
+    if stack_centered is not None and stack_norms is not None:
+        seed_vec = stack_centered[seed_ix, seed_iz, :]
+        seed_norm = float(stack_norms[seed_ix, seed_iz])
+        if seed_norm > 0:
+            dots = np.einsum('ijk,k->ij', stack_centered, seed_vec)
+            with np.errstate(divide='ignore', invalid='ignore'):
+                seed_corr_map = np.where(stack_norms > 0, dots / (stack_norms * seed_norm), 0.0)
+        else:
+            seed_corr_map = np.zeros((nx, nz))
+    else:
+        seed_corr_map = None
+        # Precompute normalised seed trace for per-neighbour correlation
+        seed_trace = stack[seed_ix, seed_iz, :].astype(float)
+        seed_centered = seed_trace - seed_trace.mean()
+        seed_norm = float(np.sqrt((seed_centered ** 2).sum()))
+
+    # Max-heap via negated correlation (heapq is a min-heap)
+    in_heap[seed_ix, seed_iz] = True
+    heap = [(-1.0, seed_ix, seed_iz)]
+
     roi_pixels = []
     min_x, max_x = seed_ix, seed_ix
     min_z, max_z = seed_iz, seed_iz
 
     reached_x_max = False
     reached_z_max = False
 
-    while queue and len(roi_pixels) < n_pix_max:
-        ix, iz, corr_i = queue.pop()
+    while heap and len(roi_pixels) < n_pix_max:
+        neg_corr_i, ix, iz = heapq.heappop(heap)
+        corr_i = -neg_corr_i
+
         if visited[ix, iz]:
+            # Pixel was invalidated by a spatial-bound update after being pushed
             continue
 
         visited[ix, iz] = True
@@ -466,32 +518,41 @@ def corr_map_grow_roi(
 
             # update range
             if not reached_x_max:
-                min_x = np.minimum(ix, min_x)
-                max_x = np.maximum(ix, max_x)
+                min_x = min(ix, min_x)
+                max_x = max(ix, max_x)
 
                 if (max_x - min_x) >= n_pix_max_x - 1:
                     visited[:min_x, :] = True
                     visited[max_x + 1:, :] = True
+                    in_heap[:min_x, :] = True
+                    in_heap[max_x + 1:, :] = True
                     reached_x_max = True
 
             if not reached_z_max:
-                min_z = np.minimum(iz, min_z)
-                max_z = np.maximum(iz, max_z)
+                min_z = min(iz, min_z)
+                max_z = max(iz, max_z)
 
                 if (max_z - min_z) >= n_pix_max_z - 1:
                     visited[:, :min_z] = True
                     visited[:, max_z + 1:] = True
+                    in_heap[:, :min_z] = True
+                    in_heap[:, max_z + 1:] = True
                     reached_z_max = True
 
             for dx, dz in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
-                new_ix, new_iy = ix + dx, iz + dz
-                if 0 <= new_ix < nx and 0 <= new_iy < nz and not visited[new_ix, new_iy]:
-                    if corr_map is not None:
-                        new_corr_i = corr_map[new_ix, new_iy]
+                new_ix, new_iz = ix + dx, iz + dz
+                if 0 <= new_ix < nx and 0 <= new_iz < nz and not in_heap[new_ix, new_iz]:
+                    in_heap[new_ix, new_iz] = True
+                    if seed_corr_map is not None:
+                        new_corr_i = seed_corr_map[new_ix, new_iz]
                     else:
-                        new_corr_i = np.corrcoef(stack[seed_ix, seed_iz, :], stack[new_ix, new_iy, :])[0, 1]
-                    queue.append((new_ix, new_iy, new_corr_i))
-
-        queue = deque(sorted(queue, key=lambda v: v[2]))
+                        neigh_trace = stack[new_ix, new_iz, :].astype(float)
+                        neigh_centered = neigh_trace - neigh_trace.mean()
+                        neigh_norm = float(np.sqrt((neigh_centered ** 2).sum()))
+                        if seed_norm > 0 and neigh_norm > 0:
+                            new_corr_i = float(np.dot(seed_centered, neigh_centered) / (seed_norm * neigh_norm))
+                        else:
+                            new_corr_i = 0.0
+                    heapq.heappush(heap, (-new_corr_i, new_ix, new_iz))
 
     return roi_pixels