Focusing with line scan

src/eaa/image_proc.py

@@ -25,4 +25,42 @@ def stitch_images(images: list[np.ndarray], gap: int = 0) -> np.ndarray:
         buffer[:img.shape[0], x : x + img.shape[1]] = img
         x += img.shape[1] + gap
     return buffer
-
+
+
+def windowed_phase_cross_correlation(
+    moving: np.ndarray, 
+    ref: np.ndarray, 
+) -> np.ndarray:
+    """Phase correlation with windowing.
+
+    Parameters
+    ----------
+    moving : np.ndarray
+        A 2D image.
+    ref : np.ndarray
+        A 2D image.
+
+    Returns
+    -------
+    np.ndarray
+        The shift of the moving image with respect to the reference image.
+    """
+    assert np.all(np.array(moving.shape) == np.array(ref.shape)), (
+        "The shapes of the moving and reference images must be the same."
+    )
+    win_y = np.hanning(moving.shape[0])
+    win_x = np.hanning(moving.shape[1])
+    win = np.outer(win_y, win_x)
+
+    f_moving = np.fft.fft2(moving * win)
+    f_ref = np.fft.fft2(ref * win)
+
+    f_corr = f_moving * f_ref.conj()
+    f_corr = f_corr / np.abs(f_corr)
+
+    map = np.fft.ifft2(f_corr).real
+    shift = np.array(np.unravel_index(np.argmax(map), map.shape))
+    for i in range(2):
+        if shift[i] > map.shape[i] / 2:
+            shift[i] -= map.shape[i]
+    return shift

src/eaa/maths.py

@@ -29,7 +29,7 @@ def gaussian_1d(x: np.ndarray, a: float, mu: float, sigma: float, c: float = 0)
 def fit_gaussian_1d(
     x: np.ndarray,
     y: np.ndarray,
-    y_threshold: float = 0.3,
+    y_threshold: float = 0,
 ) -> tuple[float, float, float]:
     """Fit a 1D Gaussian to the data after subtracting a linear background.
 

src/eaa/task_managers/base.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Callable
 import sqlite3
 import logging
 import time
@@ -327,6 +327,8 @@ def run_feedback_loop(
         n_first_images_to_keep: Optional[int] = None,
         n_past_images_to_keep: Optional[int] = None,
         allow_non_image_tool_responses: bool = True,
+        hook_functions: Optional[dict[str, Callable]] = None,
+        *args, **kwargs
     ) -> None:
         """Run an agent-involving feedback loop.
 
@@ -364,7 +366,20 @@ def run_feedback_loop(
         allow_non_image_tool_responses : bool, optional
             If False, the agent will be asked to redo the tool call if it returns
             anything that is not an image path.
+        hook_functions : dict[str, Callable], optional
+            A dictionary of hook functions to call at certain points in the loop.
+            The keys specify the points where the hook functions are called, and
+            the values are the callables. Allowed keys are:
+            - `image_path_tool_response`: 
+              args: {"img_path": str}
+              return: {"response": Dict[str, Any], "outgoing": Dict[str, Any]}
+              Executed when the tool response is an image path, after the tool
+              response is added to the context but before the image is loaded and
+              sent to the agent. When this function is given, it **replaces** the
+              `agent.receive` call so be sure to send the image to the agent in
+              the hook if this is intended.
         """
+        hook_functions = hook_functions or {}
         round = 0
         image_path = None
         response, outgoing = self.agent.receive(
@@ -404,12 +419,15 @@ def run_feedback_loop(
 
                 if tool_response_type == ToolReturnType.IMAGE_PATH:
                     image_path = tool_response["content"]
-                    response, outgoing = self.agent.receive(
-                        message_with_acquired_image,
-                        image_path=image_path,
-                        context=self.context,
-                        return_outgoing_message=True
-                    )
+                    if "image_path_tool_response" in hook_functions:
+                        response, outgoing = hook_functions["image_path_tool_response"](image_path)
+                    else:
+                        response, outgoing = self.agent.receive(
+                            message_with_acquired_image,
+                            image_path=image_path,
+                            context=self.context,
+                            return_outgoing_message=True
+                        )
                 elif tool_response_type == ToolReturnType.EXCEPTION:
                     response, outgoing = self.agent.receive(
                         "The tool returned an exception. Please fix the exception and try again.",