From 817f244d62526080d69fab357ab5f2b4240b98c2 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 6 Feb 2026 15:56:53 +0100 Subject: [PATCH] Added a script to test autofocus, WIP --- src/aare/gui/auto_focus_test.py | 670 ++++++++++++++++++++++++++++++++ 1 file changed, 670 insertions(+) create mode 100644 src/aare/gui/auto_focus_test.py diff --git a/src/aare/gui/auto_focus_test.py b/src/aare/gui/auto_focus_test.py new file mode 100644 index 00000000..39cd3f98 --- /dev/null +++ b/src/aare/gui/auto_focus_test.py @@ -0,0 +1,670 @@ +import time +from typing import Callable, Iterable + +import cv2 +import numpy as np +from aare.common.autofocus_tools import focus_measure_edges +from aare.common.beamline import mx_beamline +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.models import AutofocusSettings +from aare.daq.config import BeamlineConfig +from aare.daq.daq import AareDAQ +from aare.daq.devices import BeamlineDevices + +def make_circular_mask(shape_hw: tuple[int, int], center_x: float, center_y: float, radius: float) -> np.ndarray: + h, w = int(shape_hw[0]), int(shape_hw[1]) + y, x = np.ogrid[:h, :w] + return (x - float(center_x)) ** 2 + (y - float(center_y)) ** 2 <= float(radius) ** 2 + +def _parabola_vertex(x1, y1, x2, y2, x3, y3) -> float | None: + # Fit parabola through 3 points; return vertex x if it's a maximum. + denom = (x1 - x2) * (x1 - x3) * (x2 - x3) + if abs(denom) < 1e-15: + return None + a = (x3 * (y2 - y1) + x2 * (y1 - y3) + x1 * (y3 - y2)) / denom + b = (x3**2 * (y1 - y2) + x2**2 * (y3 - y1) + x1**2 * (y2 - y3)) / denom + if a >= 0: + return None + return float(-b / (2 * a)) + +class AutofocusController: + """ + Fast autofocus: bracket -> ternary -> optional parabola. + + You inject: + - get_gray_image(): np.ndarray (2D) + - get_frame_id(): int (UniqueId) OR None + - move_to(z): move stage to requested z (units are up to you) + - wait_for_stop(): block until motion ends + + The key speed/robustness trick is waiting for a *new frame id* after motion. + """ + def __init__( + self, + *, + get_gray_image, + focus_measure, + move_to, + wait_for_stop, + get_frame_id=None, + fps: float = 25.0, + ): + self.get_gray_image = get_gray_image + self.focus_measure = focus_measure + self.move_to = move_to + self.wait_for_stop = wait_for_stop + self.get_frame_id = get_frame_id + self.fps = float(fps) + + self._uid_stuck_count = 0 + self._uid_stuck_disable_after = 3 + + def _wait_new_frames(self, frames: int = 1, timeout_s: float = 0.12) -> bool: + """ + Wait for new frames by UniqueId. + If uid appears stuck (common in standalone tests if acquisition isn't running), + quickly fall back to a short sleep so autofocus stays fast. + """ + if self.get_frame_id is None or self._uid_stuck_count >= self._uid_stuck_disable_after: + time.sleep(max(0.0, float(frames)) / max(1e-6, self.fps)) + return True + start = int(self.get_frame_id()) + print(f"Waiting for {frames} frames (uid={start})...") + target = start + int(frames) + deadline = time.perf_counter() + float(timeout_s) + + while time.perf_counter() < deadline: + if int(self.get_frame_id()) >= target: + self._uid_stuck_count = 0 + print(f"Acquired {frames} frames (uid={target}, start={start})") + return True + time.sleep(0.001) + + # uid didn't advance in time -> count as "stuck" and fall back + print(f"UID stuck for {timeout_s} s, falling back to sleep...") + self._uid_stuck_count += 1 + time.sleep(1.0 / max(1e-6, self.fps)) + return False + + def _score_at(self, z, mask: np.ndarray | None, robust_frames: int) -> float: + st = time.perf_counter() + self.move_to(z) + print(f"move_to command to z={z:.2f} (t={time.perf_counter() - st:.5f} s)") + st = time.perf_counter() + self.wait_for_stop() + print(f"wait_for_stop command (t={time.perf_counter() - st:.5f} s)") + + # Ensure next image is not a stale buffer + print("Waiting for new frame...") + st = time.perf_counter() + self._wait_new_frames(frames=1, timeout_s=0.4) + print(f"Acquired new frame (t={time.perf_counter() - st:.5f} s)") + st = time.perf_counter() + if robust_frames <= 1: + gray = self.get_gray_image() + print(f"got grey image (t={time.perf_counter() - st:.5f} s)") + return float(self.focus_measure(gray, mask)) + + vals: list[float] = [] + for _ in range(int(robust_frames)): + gray = self.get_gray_image() + vals.append(float(self.focus_measure(gray, mask))) + self._wait_new_frames(frames=1, timeout_s=0.4) + print(f'Got values after {time.perf_counter() - st:.5f} s: {vals}') + return float(np.median(np.asarray(vals, dtype=np.float64))) + + def run_once( + self, + *, + z0: float, + z_range: float, + mask: np.ndarray | None = None, + robust_frames: int = 1, + ternary_iters: int = 4, + do_parabola: bool = True, + edge_stop: bool = True, + flat_rel_tol: float = 0.03, + ) -> tuple[float, float]: + """ + Returns (best_z, best_focus). + + edge_stop: + If True and the best bracket point is at ±0.5*z_range, stop early. + (Means the peak is likely outside the search window.) + + flat_rel_tol: + If (max-min)/max is below this, treat focus curve as flat and stop early. + """ + R = float(z_range) + + # 1) 5-point bracket + zs = np.array( + [z0 - 0.5 * R, z0 - 0.25 * R, z0, z0 + 0.25 * R, z0 + 0.5 * R], + dtype=np.float64, + ) + fs = np.array([self._score_at(float(z), mask, robust_frames) for z in zs], dtype=np.float64) + f0 = float(fs[2]) # z0 + f_max = float(fs.max()) + if f0 > 0 and (f_max / f0) < 1.05: # <5% improvement available + return float(zs[2]), f0 + + best_i = int(np.argmax(fs)) + z_best = float(zs[best_i]) + f_best = float(fs[best_i]) + + # Early exit if the curve is basically flat (no meaningful improvement) + f_min = float(fs.min()) + if f_max > 0 and (f_max - f_min) / f_max < float(flat_rel_tol): + return z_best, f_best + + # Early exit if best is at range edge: bracket does not contain a maximum + if edge_stop and (best_i == 0 or best_i == len(zs) - 1): + return z_best, f_best + + # Local bracket for ternary search + iL = max(0, best_i - 1) + iR = min(len(zs) - 1, best_i + 1) + zL, zR = float(zs[iL]), float(zs[iR]) + + sampled: dict[float, float] = {float(zs[i]): float(fs[i]) for i in range(len(zs))} + + if zL == zR: + return z_best, f_best + + # 2) ternary search in local bracket (assumes unimodal-ish) + for _ in range(int(ternary_iters)): + a, b = (zL, zR) if zL < zR else (zR, zL) + z1 = a + (b - a) / 3.0 + z2 = b - (b - a) / 3.0 + + if z1 not in sampled: + sampled[z1] = self._score_at(float(z1), mask, robust_frames) + if z2 not in sampled: + sampled[z2] = self._score_at(float(z2), mask, robust_frames) + + if sampled[z1] < sampled[z2]: + zL = z1 + else: + zR = z2 + + # 3) optional 3-point parabola around current best sample + if do_parabola and len(sampled) >= 3: + items = sorted(sampled.items(), key=lambda t: t[0]) + zz = np.array([p[0] for p in items], dtype=np.float64) + ff = np.array([p[1] for p in items], dtype=np.float64) + k = int(np.argmax(ff)) + + if 0 < k < len(zz) - 1: + zv = _parabola_vertex( + float(zz[k - 1]), float(ff[k - 1]), + float(zz[k]), float(ff[k]), + float(zz[k + 1]), float(ff[k + 1]), + ) + if zv is not None and float(zz[k - 1]) <= zv <= float(zz[k + 1]): + if zv not in sampled: + sampled[zv] = self._score_at(float(zv), mask, robust_frames) + + z_best, f_best = max(sampled.items(), key=lambda t: t[1]) + return float(z_best), float(f_best) + +def __auto_focus(settings: AutofocusSettings) -> float: + """ + Fast autofocus on Smargon Z: + - bracket (5 points) + - ternary search (few iters) + - optional parabola refine + + Returns: + Best Z offset in mm (beamline Z delta) relative to the starting position. + """ + geom = daq.sample_geometry + start_smargon = devs.smargon_pos + + # ROI center: use provided, else use beam location (beam mark) + center_x = float(settings.center_x_pxl) if settings.center_x_pxl is not None else float(geom.beam_location_pxl.x) + center_y = float(settings.center_y_pxl) if settings.center_y_pxl is not None else float(geom.beam_location_pxl.y) + radius_pxl = float(settings.radius_pxl) + + z_range_mm = float(settings.z_range_um) / 1000.0 + z_steps = int(settings.z_steps) + + # Build mask once (needs image shape) + first = daq.camera_image_gray + if first is None: + raise RuntimeError("Autofocus: no camera image available.") + if first.ndim != 2: + raise RuntimeError("Autofocus: expected grayscale image (2D).") + + #mask = make_circular_mask(first.shape[:2], center_x=center_x, center_y=center_y, radius=radius_pxl) + + height, width = first.shape + y, x = np.ogrid[:height, :width] + mask = (x - center_x) ** 2 + (y - center_y) ** 2 <= radius_pxl ** 2 + + def move_to_delta_z_mm(dz_mm: float) -> None: + # Apply relative motion in *beamline Z* via the geometry transform + sh_new = start_smargon.sh_mm + geom.smargon_nudge(Coordinate(z=float(dz_mm))) + target = SmargonCoordinate( + sh_mm=sh_new, + phi_deg=start_smargon.phi_deg, + chi_deg=start_smargon.chi_deg, + ) + devs.smargon_pos = target + + def wait_for_stop() -> None: + devs.smargon_wait(timeout=30) + + def get_gray() -> np.ndarray: + img = daq.camera_image_gray + if img is None: + raise RuntimeError("Autofocus: failed to acquire image.") + return img + + ctrl = StepwiseAutofocus( + get_gray_image=get_gray, + focus_measure=focus_measure_edges, + move_to=move_to_delta_z_mm, + wait_for_stop=wait_for_stop, + get_frame_id=devs.samcam_frame_id, + fps=25.0, + ) + + # Robustness vs speed: + # - 1 is fastest + # - 2 is more stable (median of 2 frames) and often still < 1 s total + robust_frames = 1 + + best_dz, best_f, zs, fs = ctrl.run(z0=0.0, z_range=z_range_mm, z_steps=z_steps, mask=mask, + refine=False, include_baseline=False) + + # Move to the best position (controller ends at last probed z; ensure final is best) + move_to_delta_z_mm(best_dz) + wait_for_stop() + + print( + f"Autofocus complete: best_dz={best_dz * 1000.0:.1f} um, focus={best_f:.2f}, " + f"roi_center=({center_x:.1f},{center_y:.1f}), r={radius_pxl:.1f}px" + ) + + return float(best_dz) + + # def auto_focus(self, settings: AutofocusSettings) -> float: + # """ + # Public autofocus method. Only allowed in SampleAlignment state. + # Returns best Z offset in mm (beamline Z delta) relative to start. + # """ + # self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) + # try: + # best_dz_mm = self.__auto_focus(settings) + # self.__cfg.state_busy = False + # return best_dz_mm + # except Exception as e: + # logger.error(f"Autofocus failed: {e}") + # self.__cfg.state_busy = False + # raise +class StepwiseAutofocus: + def __init__(self, *, get_gray_image, focus_measure, move_to, wait_for_stop, get_frame_id=None, fps=25.0): + self.get_gray_image = get_gray_image + self.focus_measure = focus_measure + self.move_to = move_to + self.wait_for_stop = wait_for_stop + self.get_frame_id = get_frame_id + self.fps = float(fps) + + def _wait_new_frame(self, timeout_s: float = 0.25) -> None: + if self.get_frame_id is None: + time.sleep(1.0 / max(1e-6, self.fps)) + return + start = int(self.get_frame_id()) + deadline = time.perf_counter() + float(timeout_s) + while time.perf_counter() < deadline: + if int(self.get_frame_id()) > start: + return + time.sleep(0.001) + # fallback: don't hang + time.sleep(1.0 / max(1e-6, self.fps)) + + def score_at(self, z: float, mask=None) -> float: + self.move_to(float(z)) + self.wait_for_stop() + self._wait_new_frame(timeout_s=0.25) + gray = self.get_gray_image() + return float(self.focus_measure(gray, mask)) + + def run( + self, + *, + z0: float, + z_range: float, + z_steps: int, + mask=None, + refine: bool = False, + include_baseline: bool = False, + ) -> tuple[float, float, np.ndarray, np.ndarray]: + """ + Returns: + (best_z, best_focus, z_positions, focus_values) + + If include_baseline=False and refine=False, this will evaluate focus exactly `z_steps` times. + """ + z_steps = int(z_steps) + if z_steps < 3: + raise ValueError("z_steps must be >= 3 for a meaningful scan.") + + if include_baseline: + _ = self.score_at(float(z0), mask=mask) + + zs = np.linspace(z0 - 0.5 * float(z_range), z0 + 0.5 * float(z_range), z_steps, dtype=np.float64) + fs = np.empty_like(zs) + + for i, z in enumerate(zs): + fs[i] = self.score_at(float(z), mask=mask) + + best_i = int(np.argmax(fs)) + best_z = float(zs[best_i]) + best_f = float(fs[best_i]) + + if refine and 0 < best_i < (len(zs) - 1): + dz = float(zs[best_i + 1] - zs[best_i]) + z_candidates = np.array([best_z - dz, best_z, best_z + dz], dtype=np.float64) + f_candidates = np.array([self.score_at(float(zc), mask=mask) for zc in z_candidates], dtype=np.float64) + j = int(np.argmax(f_candidates)) + best_z = float(z_candidates[j]) + best_f = float(f_candidates[j]) + + return best_z, best_f, zs, fs + +def __auto_focus_with_aerotech(settings: AutofocusSettings) -> float: + """ + 1) Fast focus scan on Aerotech GMZ (true focus axis) + 2) Return GMZ to home position + 3) Apply one Smargon move to preserve the focus (using a local Jacobian estimate) + + Returns: + Smargon delta (in the same "beamline z command" units you use in geom.smargon_nudge(Coordinate(z=...))). + """ + geom = daq.sample_geometry + start_smargon = devs.smargon_pos + + center_x = float(geom.beam_location_pxl.x) + center_y = float(geom.beam_location_pxl.y) + radius_pxl = float(settings.radius_pxl) + + z_range_mm = float(settings.z_range_um) / 1000.0 + z_steps = int(settings.z_steps) + + def get_gray() -> np.ndarray: + """ + Match GUI pipeline: + - if Bayer: debayer -> RGB + - flip horizontally + - convert to gray (uint8) + """ + img = daq.camera_image # <-- NOTE: use raw, not camera_image_gray + if img is None: + raise RuntimeError("Autofocus: failed to acquire image.") + + # If already grayscale + if img.ndim == 2: + bayer = img.astype(np.uint8, copy=False) + rgb = cv2.cvtColor(bayer, cv2.COLOR_BAYER_GB2RGB) + rgb = rgb[:, ::-1, :].copy() + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + return gray + + # If RGB-like + if img.ndim == 3 and img.shape[2] >= 3: + rgb = img[:, :, :3] + rgb = rgb[:, ::-1, :].copy() + if rgb.dtype != np.uint8: + rgb = np.clip(rgb, 0, 255).astype(np.uint8) + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + return gray + + raise RuntimeError(f"Autofocus: unexpected image shape {img.shape}") + + + first = get_gray() + print(first.shape[:2]) + print(center_x, center_y, radius_pxl) + print((first.shape[1]-1) - center_x) + if first is None or first.ndim != 2: + raise RuntimeError("Autofocus: no grayscale image available.") + mask = make_circular_mask(first.shape[:2], center_x=center_x, center_y=center_y, radius=radius_pxl) + + def score_focus() -> float: + gray = get_gray() + + # Ensure we're comparing apples-to-apples in logs + g = gray + if g.dtype != np.uint8: + g_u8 = np.clip(g, 0, 255).astype(np.uint8) + else: + g_u8 = g + + roi = g_u8[mask] + mean_dn = float(roi.mean()) if roi.size else 0.0 + std_dn = float(roi.std()) if roi.size else 0.0 + + raw = float(focus_measure_edges(g_u8, mask)) + + # Normalize to reduce exposure/gain dependence (gradient energy scales ~ intensity^2) + norm = raw / ((mean_dn + 1e-6) ** 2) + + print(f"AF: mean={mean_dn:.1f} std={std_dn:.1f} raw_focus={raw:.2f} norm_focus={norm:.6f}") + return norm + + # --------- + # A) Aerotech GMZ scan (relative to current GMZ = "home" for this autofocus call) + # --------- + aero0 = devs.aerotech_pos + gmz0 = float(aero0.z) + + gmz_offsets = np.linspace(-0.5 * z_range_mm, 0.5 * z_range_mm, z_steps, dtype=np.float64) + gmz_scores = [] + + for dz in gmz_offsets: + devs.aerotech.move_motor_linear("Z", gmz0 + float(dz), 10) + # wait 1 new frame after motion so we don't score an old buffer + start_uid = devs.samcam_frame_id() + t_deadline = time.perf_counter() + 0.25 + while time.perf_counter() < t_deadline and devs.samcam_frame_id() == start_uid: + time.sleep(0.001) + gmz_scores.append(score_focus()) + + gmz_scores = np.asarray(gmz_scores, dtype=np.float64) + best_i = int(np.argmax(gmz_scores)) + best_gmz_offset = float(gmz_offsets[best_i]) + best_focus = float(gmz_scores[best_i]) + + # Move GMZ back to "home" (gmz0) + devs.aerotech.move_motor_absolute("Z", gmz0, 10000) + + # If best was ~0 anyway, nothing to bake in + if abs(best_gmz_offset) < 1e-6: + print(f"Aerotech prefocus: best_gmz_offset≈0, focus={best_focus:.2f}") + return 0.0 + + # --------- + # B) Estimate local Jacobian: how Aerotech GMZ changes per unit Smargon beamline-z command + # We do two probe moves in the Smargon command space and measure GMZ readback. + # --------- + def move_smargon_beamline_dz(dz_mm: float) -> None: + sh_new = start_smargon.sh_mm + geom.smargon_nudge(Coordinate(z=float(dz_mm))) + target = SmargonCoordinate( + sh_mm=sh_new, + phi_deg=start_smargon.phi_deg, + chi_deg=start_smargon.chi_deg, + ) + devs.smargon_pos = target + devs.smargon_wait(timeout=30) + + move_smargon_beamline_dz(best_gmz_offset) + + print( + f"Aerotech prefocus: best_gmz_offset={best_gmz_offset*1000} um, focus={best_focus:.2f} " + f"Smargon_start: {start_smargon.sh_mm} um, Smargon_end: {devs.smargon_pos.sh_mm} um" + ) + return float(best_gmz_offset) + +def focus_measure_laplacian(gray: np.ndarray, mask: np.ndarray | None = None) -> float: + """ + Fast focus metric: variance of Laplacian. + + Notes: + - Works best on uint8 images. + - Use a mask/ROI to avoid scoring irrelevant background. + """ + if gray is None: + return 0.0 + if gray.ndim != 2: + raise ValueError(f"Expected 2D grayscale image, got shape={gray.shape}") + + g = gray + if g.dtype != np.uint8: + g = np.clip(g, 0, 255).astype(np.uint8) + + if mask is not None: + roi = g[mask] + if roi.size < 64: # too few pixels -> unstable variance + return 0.0 + # Laplacian needs 2D input; reshape ROI to a thin image is awkward. + # Better: compute Laplacian on full image and then mask the result. + lap = cv2.Laplacian(g, cv2.CV_64F, ksize=3) + v = float(lap[mask].var()) + return v + + lap = cv2.Laplacian(g, cv2.CV_64F, ksize=3) + return float(lap.var()) + + +def _wait_for_new_uid( + get_frame_id: Callable[[], int] | None, + last_uid: int | None, + *, + frames: int = 1, + timeout_s: float = 0.30, + poll_s: float = 0.002, + fallback_sleep_s: float = 0.04, +) -> int | None: + """ + Wait until UniqueId advances by `frames`. + Returns the new uid (or last_uid if we couldn't observe advancement). + """ + if get_frame_id is None: + time.sleep(fallback_sleep_s) + return last_uid + + try: + uid0 = int(get_frame_id()) if last_uid is None else int(last_uid) + except Exception: + time.sleep(fallback_sleep_s) + return last_uid + + target = uid0 + int(frames) + deadline = time.perf_counter() + float(timeout_s) + + while time.perf_counter() < deadline: + try: + uid = int(get_frame_id()) + except Exception: + uid = uid0 + + if uid >= target: + return uid + + time.sleep(poll_s) + + # Timeout: don't hang autofocus; just do a small sleep to reduce stale-buffer chance. + time.sleep(fallback_sleep_s) + return uid0 + + +def autofocus_gpt( + z_positions: Iterable[float], + move_stage_fn: Callable[[float], None], + *, + get_frame_id: Callable[[], int] | None = None, + wait_for_stop: Callable[[], None] | None = None, + mask: np.ndarray | None = None, + robust_frames: int = 1, +) -> tuple[float, list[tuple[float, float]]]: + """ + Simple autofocus scan with reliability improvements: + - waits for a new UniqueId after motion (avoids scoring stale frames) + - optional median-of-N scoring per z + """ + measures: list[tuple[float, float]] = [] + last_uid: int | None = None + + # Prime last_uid so the first point also waits for a "fresh" frame + if get_frame_id is not None: + try: + last_uid = int(get_frame_id()) + except Exception: + last_uid = None + + for z in z_positions: + move_stage_fn(float(z)) + if wait_for_stop is not None: + wait_for_stop() + + # Wait for camera to deliver a frame AFTER the move + last_uid = _wait_for_new_uid(get_frame_id, last_uid, frames=1, timeout_s=0.35) + + if robust_frames <= 1: + img = daq.camera_image_gray + score = focus_measure_laplacian(img, mask=mask) + else: + vals: list[float] = [] + for _ in range(int(robust_frames)): + img = daq.camera_image_gray + vals.append(focus_measure_laplacian(img, mask=mask)) + last_uid = _wait_for_new_uid(get_frame_id, last_uid, frames=1, timeout_s=0.35) + score = float(np.median(np.asarray(vals, dtype=np.float64))) + + measures.append((float(z), float(score))) + print(f"Z={z:.6f}, sharpness={score:.3f}") + + best_z = max(measures, key=lambda x: x[1])[0] + return best_z, measures + +# ---- Example z positions ---- +coarse = np.linspace(-0.1, 0.1, 10) # 0 to 200 microns in 10µm steps + +def move_stage(z): + # Insert your hardware code here: + print(z) + devs.aerotech.move_motor_absolute("Z", z, 1000) + #devs.aerotech.controller.read_status() + # e.g. serial.write(f"MOVE Z {z}") + + +def get_frame_id(): + return int(devs.samcam_frame_id()) + +if __name__ == "__main__": + devs = BeamlineDevices(mx_beamline()) + cfg = BeamlineConfig(mx_beamline()) + daq = AareDAQ(cfg, bl=mx_beamline()) + zoom = devs.zoom + beam_center = cfg.get_beam_mark(zoom) + settings = AutofocusSettings(center_x_pxl=beam_center[0], center_y_pxl=beam_center[1], + radius_pxl=30, z_range_um=400, z_steps=40) + st = time.perf_counter() + best_z, curve = autofocus_gpt(coarse, move_stage, get_frame_id=get_frame_id) + move_stage(0) + #move_stage(best_z) + geom = daq.sample_geometry + start_smargon = devs.smargon_pos + + sh_new = start_smargon.sh_mm + geom.smargon_nudge(Coordinate(z=float(best_z))) + target = SmargonCoordinate( + sh_mm=sh_new, + phi_deg=start_smargon.phi_deg, + chi_deg=start_smargon.chi_deg, + ) + devs.smargon_pos = target + devs.smargon_wait(timeout=30) + print("Best focus at:", best_z) + print(f"Total time: {time.perf_counter() - st:.5f} s") \ No newline at end of file