From 8c45eb7be9046cf0416cb0dbd454329bd84d9671 Mon Sep 17 00:00:00 2001 From: leonarski_f Date: Fri, 16 Jan 2026 17:14:54 +0100 Subject: [PATCH] GUI: Autofocus - work in progress --- src/aare/common/models.py | 4 +- src/aare/daq/daq.py | 160 +++++++++++++++++++++++++- src/aare/daq/server.py | 2 +- src/aare/gui/threads/camera_thread.py | 5 +- src/aare/gui/widgets/camera_image.py | 7 +- 5 files changed, 169 insertions(+), 9 deletions(-) diff --git a/src/aare/common/models.py b/src/aare/common/models.py index fd7827ee..da028d0a 100644 --- a/src/aare/common/models.py +++ b/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 diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 387aba08..a1745fed 100644 --- a/src/aare/daq/daq.py +++ b/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 diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index b46c2316..f482271b 100644 --- a/src/aare/daq/server.py +++ b/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") diff --git a/src/aare/gui/threads/camera_thread.py b/src/aare/gui/threads/camera_thread.py index a2b9d4d6..15214f48 100644 --- a/src/aare/gui/threads/camera_thread.py +++ b/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) diff --git a/src/aare/gui/widgets/camera_image.py b/src/aare/gui/widgets/camera_image.py index 19e54aef..eefdcf24 100644 --- a/src/aare/gui/widgets/camera_image.py +++ b/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())