From 7d65a18d28565b2402a1240d296d39ec6881400e Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 30 May 2025 16:29:16 +0200 Subject: [PATCH] Autofocus function --- common/src/aaredaqlib/models.py | 7 ++++++ daq/src/aaredaq/autofocus.py | 31 +++++++++++++++++++++++++ daq/src/aaredaq/daq.py | 29 +++++++++++++++++++++-- daq/src/aaredaq/devices.py | 1 - daq/src/aaredaq/server.py | 10 ++++++-- daq/src/mxlibs3/area_detector.py | 23 ++++++++++++++++++ gui/src/aaregui/main_window.py | 2 ++ gui/src/aaregui/threads/daq_worker.py | 7 +++++- gui/src/aaregui/widgets/camera_image.py | 10 +++++++- 9 files changed, 113 insertions(+), 7 deletions(-) create mode 100644 daq/src/aaredaq/autofocus.py diff --git a/common/src/aaredaqlib/models.py b/common/src/aaredaqlib/models.py index e4c2ac22..d3f42a25 100644 --- a/common/src/aaredaqlib/models.py +++ b/common/src/aaredaqlib/models.py @@ -137,6 +137,13 @@ class SampleCameraSettings(BaseModel): gain: float exposure: float +class AutofocusSettings(BaseModel): + center_x_pxl: float + center_y_pxl: float + radius_pxl: float + z_range_um: float + z_steps: int + class BeamlineStatus(BaseModel): energy_keV: float dtz_mm: float diff --git a/daq/src/aaredaq/autofocus.py b/daq/src/aaredaq/autofocus.py new file mode 100644 index 00000000..cceaf17f --- /dev/null +++ b/daq/src/aaredaq/autofocus.py @@ -0,0 +1,31 @@ +import cv2 +import numpy as np + + +def calculate_focus_measure(image: np.ndarray, center_x: float, center_y: float, radius: float) -> float: + """ + Calculate the focus measure within a circular region around a point. + + Args: + image: Input image + center_x: X coordinate of the center point + center_y: Y coordinate of the center point + radius: Radius of the circular region to analyze + Returns: + Focus measure value (higher means more in focus) + """ + if len(image.shape) > 2: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + else: + gray = image + + # Create a circular mask + height, width = gray.shape + y, x = np.ogrid[:height, :width] + mask = ((x - center_x) ** 2 + (y - center_y) ** 2 <= radius ** 2) + + # Apply Laplacian operator + laplacian = cv2.Laplacian(gray, cv2.CV_64F) + + # Apply mask and calculate variance of only the circular region + return laplacian[mask].var() diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index 3b463c50..c17b8abb 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -5,12 +5,14 @@ from typing import List, Tuple import cv2 import numpy as np import redis + +from aaredaq.autofocus import calculate_focus_measure from aaredaqlib.coordinate import Coordinate, SmargonCoordinate from aaredaqlib.models import ( SampleShortInfo, PuckLoadedInfo, SampleShortInfoList, - DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, ) + DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, ) from aaredaqlib.raster_grid import RasterGridRequest from aaredaqlib.rotation_scan import RotationScanRequest from aaredaqlib.sample_geometry import SampleGeometryModel @@ -117,6 +119,28 @@ class AareDAQ: def samcam_settings(self, s: SampleCameraSettings): self.__devs.samcam_settings = s + def autofocus(self, f: AutofocusSettings): + pos = self.__cfg.abr_meas_pos + + z_min = pos.z - f.z_range_um / 2.0 + z_max = pos.z + f.z_range_um / 2.0 + + measures = [] + + for i in range(f.z_steps): + pos.z = z_min + i * (z_max - z_min) / f.z_steps + self.__devs.abr_pos = pos + self.__devs.sample_cam.get_single_image() + image = self.__devs.sample_cam.get_image(gray=True) + val = calculate_focus_measure(image, f.center_x_pxl, f.center_y_pxl, f.radius_pxl) + measures.append((pos.z, val)) + best_pos, _ = max(measures, key=lambda x: x[1]) + + pos.z = best_pos + self.__devs.abr_pos = pos + self.__cfg.abr_meas_pos = pos + self.__devs.sample_cam.collect_auto() + def tweak_abr_meas_pos(self, c: Coordinate): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: @@ -517,7 +541,8 @@ class AareDAQ: @property def sample_spreadsheet(self) -> SampleShortInfoList: - return self.__devs.aare.get_sample_info() + return SampleShortInfoList(s=[]) + #return self.__devs.aare.get_sample_info() def sample_spreadsheet_user(self, pgroup: str) -> SampleShortInfoList: sample = self.sample_spreadsheet diff --git a/daq/src/aaredaq/devices.py b/daq/src/aaredaq/devices.py index e2d28186..b30dfe47 100644 --- a/daq/src/aaredaq/devices.py +++ b/daq/src/aaredaq/devices.py @@ -198,7 +198,6 @@ class BeamlineDevices: @abr_pos.setter def abr_pos(self, pos: Coordinate): - print(f"Going to x {pos.x} y {pos.y} z {pos.z}") self.gmx.move(pos.x, wait=True) self.gmy.move(pos.y, wait=True) self.gmz.move(pos.z, wait=True) diff --git a/daq/src/aaredaq/server.py b/daq/src/aaredaq/server.py index fefbc657..7f773484 100644 --- a/daq/src/aaredaq/server.py +++ b/daq/src/aaredaq/server.py @@ -6,7 +6,7 @@ import urllib3 import uvicorn from aaredaqlib.coordinate import SmargonCoordinate, Coordinate from aaredaqlib.models import SampleShortInfo, DAQStatusModel, BeamlineStateEnum, BeamlineSettingsModel, \ - SampleShortInfoList, SessionStatus, SampleCameraSettings + SampleShortInfoList, SessionStatus, SampleCameraSettings, AutofocusSettings from aaredaqlib.raster_grid import RasterGridRequest from aaredaqlib.rotation_scan import RotationScanRequest from aaredaqlib.sample_geometry import SampleGeometryModel @@ -113,10 +113,16 @@ async def clear_beam_center(token: str = Depends(oauth2_scheme)): @app.put("/beamline/samcam") async def samcam_settings(s: SampleCameraSettings, token: str = Depends(oauth2_scheme)): - auth.check_jwt_staff(auth.parse_token(token)) + auth.check_jwt_rw(cfg, auth.parse_token(token)) daq.samcam_settings = s return "OK" +@app.post("/samcam/autofocus") +async def samcam_autofocus(s: AutofocusSettings, token: str = Depends(oauth2_scheme)): + auth.check_jwt_rw(cfg, auth.parse_token(token)) + daq.autofocus(s) + return "OK" + @app.post("/shutter/close") async def close_shutter(token: str = Depends(oauth2_scheme)): auth.check_jwt_rw(cfg, auth.parse_token(token)) diff --git a/daq/src/mxlibs3/area_detector.py b/daq/src/mxlibs3/area_detector.py index 1c2de069..082f8cb1 100644 --- a/daq/src/mxlibs3/area_detector.py +++ b/daq/src/mxlibs3/area_detector.py @@ -7,6 +7,7 @@ class epicsAD(object): self.img = None self.monitored = False + self.image_mode = epics.PV(prefix + cam + "ImageMode") self.gain = epics.PV(prefix + cam + "Gain") self.acquire = epics.PV(prefix + cam + "Acquire") self.color = epics.PV(prefix + cam + "ColorMode") @@ -22,6 +23,8 @@ class epicsAD(object): self.dim2 = epics.PV(prefix + image + "ArraySize2_RBV") self.uid = epics.PV(prefix + image + "UniqueId_RBV") self.data = epics.PV(prefix + image + "ArrayData") + self.counter = epics.PV(prefix + cam + "ArrayCounter_RBV") + self.busy = epics.PV(prefix + cam + "AcquireBusy", auto_monitor=True) try: epics.ca.pend_io() @@ -67,6 +70,26 @@ class epicsAD(object): shape = self.getShape() return (shape[0] / 2, shape[1] / 2) + def get_single_image(self): + self.image_mode.put(0) + counter = self.counter.get() + self.acquire.put(1) + while self.counter.get() == counter: + epics.poll(0.01) + + def collect_one_image(self): + self.acquire.put(0) + self.image_mode.put(0, wait=True) + + self.acquire.put(1, wait=True) + while self.busy.value == 1: + epics.poll(0.01) + + def collect_auto(self): + self.acquire.put(0) + self.image_mode.put(2) + self.acquire.put(1) + def get_image(self, gray=True) -> np.ndarray: if self.monitored: return self.img diff --git a/gui/src/aaregui/main_window.py b/gui/src/aaregui/main_window.py index 931a832b..bbfa74cf 100644 --- a/gui/src/aaregui/main_window.py +++ b/gui/src/aaregui/main_window.py @@ -155,6 +155,8 @@ class MainWindow(QMainWindow): self.data_collection.helical.show_helical_start.connect(self.camera_image.show_helical_start) self.data_collection.helical.show_helical_end.connect(self.camera_image.show_helical_end) + self.camera_image.autofocus.connect(self.daq.autofocus) + self.camera_image.set_helical_start.connect(self.data_collection.helical.set_start_point) self.camera_image.set_helical_end.connect(self.data_collection.helical.set_end_point) self.job_list_panel.auto_scan.connect(self.daq.automated_scan) diff --git a/gui/src/aaregui/threads/daq_worker.py b/gui/src/aaregui/threads/daq_worker.py index 1628fd33..e60d2325 100644 --- a/gui/src/aaregui/threads/daq_worker.py +++ b/gui/src/aaregui/threads/daq_worker.py @@ -1,7 +1,8 @@ from PySide6.QtCore import Signal, QUrl, Slot, QTimer, QObject, QByteArray from PySide6.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkReply -from aaredaqlib.models import DAQStatusModel, SampleShortInfoList, SampleShortInfo, SampleCameraSettings +from aaredaqlib.models import DAQStatusModel, SampleShortInfoList, SampleShortInfo, SampleCameraSettings, \ + AutofocusSettings from aaredaqlib.coordinate import SmargonCoordinate, Coordinate from aaredaqlib.raster_grid import RasterGridRequest from aaredaqlib.rotation_scan import RotationScanRequest @@ -157,6 +158,10 @@ class DAQWorker(QObject): def samcam_settings(self, s: SampleCameraSettings): self.generic_put("beamline/samcam", s.model_dump_json()) + @Slot(AutofocusSettings) + def autofocus(self, f: AutofocusSettings): + self.generic_post("samcam/autofocus", f.model_dump_json()) + @Slot(SmargonCoordinate) def move_smargon(self, coord: SmargonCoordinate): self.generic_put("beamline/smargon", coord.model_dump_json()) diff --git a/gui/src/aaregui/widgets/camera_image.py b/gui/src/aaregui/widgets/camera_image.py index 909d4f53..e31fe4fe 100644 --- a/gui/src/aaregui/widgets/camera_image.py +++ b/gui/src/aaregui/widgets/camera_image.py @@ -21,7 +21,7 @@ from PySide6.QtWidgets import ( QFrame, ) -from aaredaqlib.models import DAQStatusModel +from aaredaqlib.models import DAQStatusModel, AutofocusSettings from aaregui.models.bookmark import SmargonBookmarkList from aaredaqlib.coordinate import Coordinate, SmargonCoordinate from aaredaqlib.raster_grid import RasterGridRequest @@ -36,6 +36,7 @@ class SampleCameraImageLabel(QGraphicsView): set_helical_start = Signal(SmargonCoordinate) set_helical_end = Signal(SmargonCoordinate) update_beam_center = Signal(float, float) + autofocus = Signal(AutofocusSettings) def __init__( self, @@ -162,6 +163,8 @@ class SampleCameraImageLabel(QGraphicsView): grab_action = menu.addAction("Grab") grab_with_overlay_action = menu.addAction("Grab with overlay") + autofocus_action = menu.addAction("Auto-focus") + beam_mark_action = menu.addAction("Mark beam center") delete_action = None @@ -190,6 +193,11 @@ class SampleCameraImageLabel(QGraphicsView): self.__screenshot_with_dialog(overlay=False) elif action == grab_with_overlay_action: 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)) elif action == beam_mark_action: c = self.mapToScene(event.pos()) self.update_beam_center.emit(c.x(), c.y())