From 9479c951067e685c59112c7fa77fadb3eec5adf1 Mon Sep 17 00:00:00 2001 From: leonarski_f Date: Fri, 16 Jan 2026 16:14:51 +0100 Subject: [PATCH] GUI: Camera thread improvements --- src/aare/common/autofocus_tools.py | 54 +++++++++++++++ src/aare/gui/main_window.py | 4 +- src/aare/gui/threads/camera_thread.py | 69 +++++++------------ src/aare/gui/threads/prediction_subscriber.py | 42 +++++++++++ 4 files changed, 122 insertions(+), 47 deletions(-) create mode 100644 src/aare/common/autofocus_tools.py create mode 100644 src/aare/gui/threads/prediction_subscriber.py diff --git a/src/aare/common/autofocus_tools.py b/src/aare/common/autofocus_tools.py new file mode 100644 index 00000000..25056e21 --- /dev/null +++ b/src/aare/common/autofocus_tools.py @@ -0,0 +1,54 @@ +import cv2 +import numpy as np + +def focus_measure_edges(gray: np.ndarray, mask: np.ndarray | None = None) -> float: + # mild denoise (optional but usually stabilizes the curve) + gray = cv2.GaussianBlur(gray, (3, 3), 0) + + gx = cv2.Scharr(gray, cv2.CV_64F, 1, 0) + gy = cv2.Scharr(gray, cv2.CV_64F, 0, 1) + g2 = gx * gx + gy * gy + + roi = g2[mask] if mask is not None else g2.reshape(-1) + if roi.size == 0: + return 0.0 + + # threshold relative to median -> knocks out noise floor + t = float(np.median(roi) * 3.0) + strong = roi[roi > t] + if strong.size == 0: + return 0.0 + + return float(strong.mean()) # higher = sharper + +def focus_measure_blob_size(gray: np.ndarray, mask: np.ndarray | None = None) -> float: + """ + Measures sharpness for a single bright blob. + Higher = sharper (smaller blob). + """ + g = gray.astype(np.float64) + + if mask is not None: + g = np.where(mask, g, 0.0) + + # Background subtraction is crucial for blob metrics + # Use a large-ish blur as background estimate (tune ksize to your scale) + bg = cv2.GaussianBlur(g, (0, 0), sigmaX=10.0, sigmaY=10.0) + s = g - bg + s[s < 0] = 0.0 + + total = float(s.sum()) + if total <= 0: + return 0.0 + + h, w = s.shape + y, x = np.mgrid[0:h, 0:w] + + cx = float((s * x).sum() / total) + cy = float((s * y).sum() / total) + + # intensity-weighted second central moment (variance) + var = float((s * ((x - cx) ** 2 + (y - cy) ** 2)).sum() / total) + + # smaller var => sharper, so invert + return float(1.0 / (var + 1e-9)) \ No newline at end of file diff --git a/src/aare/gui/main_window.py b/src/aare/gui/main_window.py index f73f5a49..4a044cd8 100644 --- a/src/aare/gui/main_window.py +++ b/src/aare/gui/main_window.py @@ -28,7 +28,8 @@ from aare.gui.scan_logic.raster_grid_manager import RasterGridManager from aare.gui.scan_logic.rotation_scan_manager import RotationScanManager from aare.gui.scan_logic.sample_mount_logic import SampleMountLogic from aare.gui.threads.axis_video_thread import VideoThread -from aare.gui.threads.camera_thread import SampleCameraThread, PredictionSubscriber +from aare.gui.threads.camera_thread import SampleCameraThread +from aare.gui.threads.prediction_subscriber import PredictionSubscriber from aare.gui.threads.daq_worker import DAQWorker from aare.gui.threads.jfjoch_viewer import JFJochDBusClient from aare.gui.widgets.camera_image import SampleCameraImageLabel @@ -339,6 +340,7 @@ class MainWindow(QMainWindow): self.daq.update.connect(self.sample_camera.update_daq_status) self.daq.update.connect(self.tell_samples.update_daq_status) self.daq.update.connect(self.ref_tools_panel.update_daq_status) + self.daq.update.connect(self.camera_thread.update_daq_status) if self.__decoded_token.staff: self.daq.update.connect(self.beamline.beam_size.update_daq_status) diff --git a/src/aare/gui/threads/camera_thread.py b/src/aare/gui/threads/camera_thread.py index f1a61707..a2b9d4d6 100644 --- a/src/aare/gui/threads/camera_thread.py +++ b/src/aare/gui/threads/camera_thread.py @@ -6,6 +6,8 @@ import zmq from PySide6.QtCore import QThread, Signal, Slot from PySide6.QtGui import QImage, QPixmap +from aare.common.autofocus_tools import focus_measure_edges +from aare.common.models import DAQStatusModel class SampleCameraThread(QThread): # Define a signal to communicate messages from the thread to the main GUI @@ -22,6 +24,19 @@ class SampleCameraThread(QThread): self.running = True self.__measure_focus = True + self.__focus_mask = None + self.__beam_x = 0 + self.__beam_y = 0 + self.__last_beam_pos = None + self.__radius = 40 + + @Slot(DAQStatusModel) + def update_daq_status(self, s: DAQStatusModel): + self.__beam_x = s.geom.beam_location_pxl.x + self.__beam_y = s.geom.beam_location_pxl.y + if (self.__beam_x, self.__beam_y) != self.__last_beam_pos: + self.__focus_mask = None # Invalidate cache + self.__last_beam_pos = (self.__beam_x, self.__beam_y) @Slot(bool) def enable_focus_measurement(self, enabled: bool = True): @@ -46,11 +61,15 @@ class SampleCameraThread(QThread): if self.__measure_focus: gray = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY) - laplacian = cv2.Laplacian(gray, cv2.CV_64F) - self.focus_measure.emit(float(laplacian.var())) + if self.__focus_mask is None or self.__focus_mask.shape != gray.shape: + height, width = gray.shape + y, x = np.ogrid[:height, :width] + self.__focus_mask = (x - self.__beam_x) ** 2 + (y - self.__beam_y) ** 2 <= self.__radius ** 2 - qimage = QImage(rgb_image.data, header_shape[1], header_shape[0], #For ethernet need header_shape[0], header_shape[1], header_shape[0]*3, - QImage.Format.Format_RGB888) + sharpness = focus_measure_edges(gray, self.__focus_mask) + self.focus_measure.emit(sharpness) + + qimage = QImage(rgb_image.data, header_shape[1], header_shape[0], QImage.Format.Format_RGB888) self.camera_image.emit(QPixmap.fromImage(qimage)) else: print("Sample camera image has wrong dimensions") @@ -67,45 +86,3 @@ class SampleCameraThread(QThread): self.__socket.close() self.quit() self.wait() - -class PredictionSubscriber(QThread): - # emits parsed JSON payload (dict with keys: time, frame_id, shape, boxes) - prediction = Signal(dict) - - def __init__(self, pred_zmq_url: str, topic: bytes = b"", parent=None): - super().__init__(parent) - self._ctx = zmq.Context() - self._sock = self._ctx.socket(zmq.SUB) - if topic: - self._sock.setsockopt(zmq.SUBSCRIBE, topic) - else: - self._sock.setsockopt(zmq.SUBSCRIBE, b"") - self._sock.connect(pred_zmq_url) - self.running = True - - def run(self): - while self.running: - try: - parts = self._sock.recv_multipart() - if not parts: - continue - # publisher sends either raw JSON or [topic, json] - payload_bytes = parts[-1] - try: - payload = json.loads(payload_bytes.decode("utf-8")) - except Exception: - continue - self.prediction.emit(payload) - except Exception as e: - print("PredictionSubscriber error:", e) - break - - def stop(self): - self.running = False - try: - self._sock.close() - except Exception: - pass - self.quit() - print("Prediction thread stopped. Exiting...") - self.wait() \ No newline at end of file diff --git a/src/aare/gui/threads/prediction_subscriber.py b/src/aare/gui/threads/prediction_subscriber.py new file mode 100644 index 00000000..4e506465 --- /dev/null +++ b/src/aare/gui/threads/prediction_subscriber.py @@ -0,0 +1,42 @@ +import json + +import zmq +from PySide6.QtCore import QThread, Signal + +class PredictionSubscriber(QThread): + # emits parsed JSON payload (dict with keys: time, frame_id, shape, boxes) + prediction = Signal(dict) + + def __init__(self, pred_zmq_url: str, topic: bytes = b"", parent=None): + super().__init__(parent) + self._ctx = zmq.Context() + self._sock = self._ctx.socket(zmq.SUB) + if topic: + self._sock.setsockopt(zmq.SUBSCRIBE, topic) + else: + self._sock.setsockopt(zmq.SUBSCRIBE, b"") + self._sock.connect(pred_zmq_url) + self.running = True + + def run(self): + while self.running: + try: + parts = self._sock.recv_multipart() + if not parts: + continue + # publisher sends either raw JSON or [topic, json] + payload_bytes = parts[-1] + try: + payload = json.loads(payload_bytes.decode("utf-8")) + except Exception: + continue + self.prediction.emit(payload) + except Exception as e: + print("PredictionSubscriber error:", e) + break + + def stop(self): + self.running = False + self._sock.close() + self.quit() + self.wait() \ No newline at end of file