GUI: Autofocus - work in progress

src/aare/common/models.py

@@ -673,8 +673,8 @@ def zoom_manager(mode: ZoomModeEnum = ZoomModeEnum.User, beamline: MXBeamline =
 
 
 class AutofocusSettings(BaseModel):
-    center_x_pxl: float
-    center_y_pxl: float
+    center_x_pxl: float | None # Use beam center
+    center_y_pxl: float | None # Use beam center
     radius_pxl: float
     z_range_um: float
     z_steps: int

src/aare/daq/daq.py

@@ -13,7 +13,7 @@ import numpy as np
 import aare.common.face_detection as fd
 from aare.daq import workflows
 from aare.daq.aaredb import AareWrapper
-
+from aare.common.autofocus_tools import focus_measure_edges
 from aare.daq.config import BeamlineConfig
 from aare.daq.config import BeamlineStateEnum
 from aare.daq.devices import BeamlineDevices
@@ -25,7 +25,7 @@ from aare.common.logger_config import setup_logger
 from aare.common.models import (
     SampleShortInfo,
     PuckLoadedInfo,
-    SampleShortInfoList,
+    SampleShortInfoList, AutofocusSettings,
     DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, ZoomModeEnum,
     SimpleScanParameters, MLBoxModel, FluorescenceSpectrumParameterModel,
     FluorescenceSpectrumOutputModel)
@@ -352,6 +352,162 @@ class AareDAQ:
     def list_loaded_pucks(self) -> List[PuckLoadedInfo]:
         return []
 
+    def __auto_focus(self, settings: AutofocusSettings, settle_time_s: float = 1.0) -> float:
+        """
+        Scan smargon Z and find the position with maximum focus measure.
+
+        Args:
+            settings: AutofocusSettings with center, radius, range, and steps
+            settle_time_s: Time to wait after each move before capturing image
+
+        Returns:
+            Best Z position (mm) found during the scan
+        """
+        geom = self.sample_geometry
+        current_smargon = self.__devs.smargon_pos
+
+        # Get center for the mask (use beam center if not specified)
+        center_x = geom.beam_location_pxl.x
+        center_y = geom.beam_location_pxl.y
+        radius_pxl = 30
+
+        # Convert z_range from um to mm
+        z_range_mm = settings.z_range_um / 1000.0
+        n_steps = settings.z_steps
+
+        # Starting Z position (current sh_mm)
+        z_start = 0.0
+        z_min = z_start - z_range_mm / 2.0
+        z_max = z_start + z_range_mm / 2.0
+        z_step = z_range_mm / (n_steps - 1) if n_steps > 1 else 0.0
+
+        # Pre-compute mask (will be created on first image)
+        focus_mask: np.ndarray | None = None
+
+        # Collect focus measures at each Z position
+        z_positions: list[float] = []
+        focus_values: list[float] = []
+
+        print(f"Starting autofocus: z_range={z_range_mm * 1000:.1f}um, steps={n_steps}, "
+                    f"center=({center_x:.1f}, {center_y:.1f}), radius={radius_pxl:.1f}px")
+
+
+        for i in range(n_steps):
+            z_pos = z_min + i * z_step
+
+            # Move to position
+            target = SmargonCoordinate(
+                chi_deg=current_smargon.chi_deg,
+                phi_deg=current_smargon.phi_deg,
+                sh_mm=geom.beamline_to_smargon(Coordinate(z=z_pos))
+            )
+            self.__devs.smargon_pos = target
+            self.__devs.smargon_wait(timeout=30)
+
+            # Wait for mechanical settling and image stabilization
+            time.sleep(settle_time_s)
+
+            # Capture image
+            img = self.camera_image
+            img = img[:, ::-1, :].copy()  # Apply horizontal flip to match GUI
+            gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+
+            print(gray.shape)
+            if gray is None:
+                logger.warning(f"Failed to get image at z={z_pos:.4f}")
+                continue
+
+            # Create mask on first valid image
+            if focus_mask is None or focus_mask.shape != gray.shape:
+                h, w = gray.shape
+                y, x = np.ogrid[:h, :w]
+                focus_mask = (x - center_x) ** 2 + (y - center_y) ** 2 <= radius_pxl ** 2
+
+            # Calculate focus measure
+            fm = focus_measure_edges(gray, focus_mask)
+
+            roi_pixels = gray[focus_mask] if focus_mask is not None else gray.ravel()[:2824]
+            print(f"Step {i}: gray mean={gray.mean():.1f}, roi mean={roi_pixels.mean():.1f}, "
+                  f"roi std={roi_pixels.std():.1f}, roi min={roi_pixels.min()}, roi max={roi_pixels.max()}")
+
+            z_positions.append(z_pos)
+            focus_values.append(fm)
+            logger.debug(f"Autofocus step {i + 1}/{n_steps}: z={z_pos:.4f}mm, focus={fm:.2f}")
+            print(f"Autofocus step {i + 1}/{n_steps}: z={z_pos:.4f}mm, focus={fm:.2f}")
+        if len(z_positions) < 3:
+            logger.error("Autofocus failed: not enough valid measurements")
+            # Return to original position
+            self.__devs.smargon_pos = current_smargon
+            self.__devs.smargon_wait(timeout=30)
+            return z_start
+
+        # Find best position - use parabolic fit around the peak for sub-step precision
+        z_arr = np.array(z_positions)
+        fm_arr = np.array(focus_values)
+
+        # Find index of maximum
+        peak_idx = int(np.argmax(fm_arr))
+
+        # Try parabolic fit if peak is not at the edge
+        if 0 < peak_idx < len(fm_arr) - 1:
+            # Fit parabola to 3 points around peak: f(z) = a*z^2 + b*z + c
+            z_fit = z_arr[peak_idx - 1: peak_idx + 2]
+            fm_fit = fm_arr[peak_idx - 1: peak_idx + 2]
+            try:
+                coeffs = np.polyfit(z_fit, fm_fit, 2)
+                a, b, c = coeffs
+                if a < 0:  # Parabola opens downward (valid peak)
+                    best_z = -b / (2 * a)
+                    # Sanity check: best_z should be within the fitted range
+                    if z_fit[0] <= best_z <= z_fit[2]:
+                        logger.info(f"Autofocus: parabolic fit found peak at z={best_z:.4f}mm")
+                    else:
+                        best_z = z_arr[peak_idx]
+                        logger.info(f"Autofocus: parabolic fit out of range, using sample peak z={best_z:.4f}mm")
+                else:
+                    best_z = z_arr[peak_idx]
+                    logger.info(f"Autofocus: invalid parabola, using sample peak z={best_z:.4f}mm")
+            except Exception as e:
+                logger.warning(f"Parabolic fit failed: {e}, using sample peak")
+                best_z = z_arr[peak_idx]
+        else:
+            best_z = z_arr[peak_idx]
+            logger.warning(f"Autofocus: peak at edge of scan range, z={best_z:.4f}mm")
+
+        # Move to best position
+        best_target = SmargonCoordinate(
+            chi_deg=current_smargon.chi_deg,
+            phi_deg=current_smargon.phi_deg,
+            sh_mm=geom.beamline_to_smargon(Coordinate(z=best_z)),
+        )
+        self.__devs.smargon_pos = best_target
+        self.__devs.smargon_wait(timeout=30)
+
+        logger.warning(f"Autofocus complete: best_z={best_z:.4f}mm, "
+                    f"focus_range=[{min(fm_arr):.2f}, {max(fm_arr):.2f}]")
+
+        return best_z
+
+    def auto_focus(self, settings: AutofocusSettings) -> float:
+        """
+        Public autofocus method. Only allowed in SampleAlignment state.
+
+        Args:
+            settings: AutofocusSettings with center, radius, range, and steps
+
+        Returns:
+            Best Z position (mm) found during the scan
+        """
+        self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment)
+        try:
+            best_z = self.__auto_focus(settings)
+            self.__cfg.state_busy = False
+            return best_z
+        except Exception as e:
+            logger.error(f"Autofocus failed: {e}")
+            self.__cfg.state_busy = False
+            raise
+
 
     def __auto_center(self, grid: RasterGridRequest) -> CompletedRasterGrid | None:
         sample = self.sample

src/aare/daq/server.py

@@ -201,7 +201,7 @@ async def samcam_settings(s: SampleCameraSettings, token: str = Depends(oauth2_s
 async def samcam_autofocus(s: AutofocusSettings, token: str = Depends(oauth2_scheme)):
     logger.debug(f"SamCam AutoFocus")
     auth.check_jwt_rw(cfg, auth.parse_token(token))
-    daq.autofocus(s)
+    daq.auto_focus(s)
     return "OK"
 
 @app.post("/beamline/shutter")

src/aare/gui/threads/camera_thread.py

@@ -58,7 +58,6 @@ class SampleCameraThread(QThread):
                     rgb_image = cv2.cvtColor(bayer_image, cv2.COLOR_BAYER_GB2RGB)
                     #cv2.COLOR_BAYER_GB2RGB for ethernet connection
                     rgb_image = rgb_image[:, ::-1, :].copy()
-
                     if self.__measure_focus:
                         gray = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY)
                         if self.__focus_mask is None or self.__focus_mask.shape != gray.shape:
@@ -67,6 +66,10 @@ class SampleCameraThread(QThread):
                             self.__focus_mask = (x - self.__beam_x) ** 2 + (y - self.__beam_y) ** 2 <= self.__radius ** 2
 
                         sharpness = focus_measure_edges(gray, self.__focus_mask)
+                        print(f"GUI: mask pixels: {self.__focus_mask.sum()}, "
+                              f"center=({self.__beam_x:.1f}, {self.__beam_y:.1f}), "
+                              f"radius={self.__radius}, "
+                              f"focus={sharpness:.2f}")
                         self.focus_measure.emit(sharpness)
 
                     qimage = QImage(rgb_image.data, header_shape[1], header_shape[0], QImage.Format.Format_RGB888)

src/aare/gui/widgets/camera_image.py

@@ -371,9 +371,10 @@ class SampleCameraImageLabel(QGraphicsView):
             self.__screenshot_with_dialog(overlay=True)
         elif action == autofocus_action:
             c = self.mapToScene(event.pos())
-            self.autofocus.emit(AutofocusSettings(center_x_pxl=c.x(), center_y_pxl=c.y(),
-                                                  radius_pxl=100, z_range_um=0.5,
-                                                  z_steps=25))
+            self.autofocus.emit(AutofocusSettings(center_x_pxl=None, center_y_pxl=None,
+                                                  radius_pxl=30,
+                                                  z_range_um=2000,
+                                                  z_steps=10))
         elif action == beam_mark_action:
             c = self.mapToScene(event.pos())
             self.update_beam_mark.emit(c.x(), c.y())