Removed old scripts
This commit is contained in:
@@ -1,670 +0,0 @@
|
||||
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")
|
||||
@@ -1,52 +0,0 @@
|
||||
from bec_lib.client import BECClient
|
||||
from bec_lib.service_config import ServiceConfig
|
||||
from bec_lib.user_macros import UserMacros
|
||||
import sys
|
||||
|
||||
sys.path.append("/sls/MX/applications/test_scripts/bec_tes")
|
||||
|
||||
service_config = ServiceConfig(redis={"host": "x06da-bec-001", "port": 6379})
|
||||
client = BECClient(service_config, name="Martins-Custom-Client")
|
||||
client.start()
|
||||
client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/calculator.py")
|
||||
client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/pxiii_parameters.py")
|
||||
client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/pxiii_energy.py")
|
||||
client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/mx_methods.py")
|
||||
client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/mx_basics.py")
|
||||
client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/energy_check.py")
|
||||
#client.macros.load_user_macro("/sls/MX/applications/test_scripts/bec_tes/y")
|
||||
dev=client.device_manager.devices
|
||||
macros = client.macros
|
||||
scans = client.scans
|
||||
|
||||
try:
|
||||
energy_ev = validate_energy(13)
|
||||
current_energy = get_current_energy()
|
||||
energy_diff = calculate_energy_difference(current_energy, energy_ev)
|
||||
dccm_pos = get_dccm_motors_positions(energy_ev)
|
||||
print(dev.dccm_theta1)
|
||||
print(dev.dccm_theta2)
|
||||
print(
|
||||
f"Moving DCCM theta1: {dccm_pos['theta1_angle']: .5g} deg, theta2: {dccm_pos['theta2_angle']: .5g} deg, "
|
||||
# f"DCM pitch: {dcm_pos['dcm_pitch']: .5g} mrad, "
|
||||
)
|
||||
theta = scans.umv(dev.dccm_theta1, dccm_pos["theta1_angle"], relative=False)
|
||||
theta.wait()
|
||||
theta_2 = scans.umv(dev.dccm_theta2, dccm_pos["theta2_angle"], relative=False)
|
||||
theta_2.wait()
|
||||
set_mirror_stripe(energy_ev)
|
||||
print(
|
||||
f"Energy difference: {energy_diff: .5g} eV, current energy: {current_energy: .5g} eV"
|
||||
)
|
||||
mono_pitch_scan(scans, plot=False)
|
||||
#bl_energy(energy_ev=13, scans=scans, plot=False)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
|
||||
# print(dev.det_z.read())
|
||||
# status = scans.mv(dev.det_z, 770.0, relative=False)
|
||||
# status.wait()
|
||||
# print(dev.det_z.read())
|
||||
client.shutdown()
|
||||
Reference in New Issue
Block a user