Autofocus function

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

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()

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

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)

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))

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

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)

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())

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())