From b75bb3746215e9946707c1ac27cce341e480dbe6 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Thu, 2 Apr 2026 14:48:40 +0200 Subject: [PATCH] DAQ: zmq_client and functionality added, please test! --- src/aare/daq/devices.py | 44 ++++++++ src/aare/daq/server.py | 13 +++ src/aare/devices/zmq_client.py | 186 +++++++++++++++++++++++++++++++++ 3 files changed, 243 insertions(+) create mode 100644 src/aare/devices/zmq_client.py diff --git a/src/aare/daq/devices.py b/src/aare/daq/devices.py index 523e5f0c..de591b3c 100644 --- a/src/aare/daq/devices.py +++ b/src/aare/daq/devices.py @@ -20,6 +20,7 @@ from aare.devices.my_motor import MyMotor from aare.devices.set_get_pv import SetGetPV, PredefinedPV from aare.devices.tell_client import make_tell_client from aare.devices.bec_worker import BECClientWorker +from aare.devices.zmq_client import ZMQCameraClient logger = setup_logger("aareDAQ") @@ -44,6 +45,12 @@ class BeamlineDevices: self.__sample_cam = epicsAD(f"{BEAMLINE}-ES-MS:") + self.__zmq_camera: ZMQCameraClient | None = None + self.__use_zmq_camera = True + if self.__use_zmq_camera: + self.__zmq_camera = ZMQCameraClient(beamline) + logger.info(f"ZMQ camera source configured: {beamline}") + self.__front_light = PredefinedPV(name='front_light', setpv=f"{BEAMLINE}-ES-FL:SET", getpv=f"{BEAMLINE}-ES-FL:SET", @@ -258,6 +265,27 @@ class BeamlineDevices: self.__sample_cam.setup(settings.gain, settings.exposure) def samcam_get_image(self, /, gray: bool = False) -> np.ndarray: + """ + Get the current sample camera image. + + Primary source: ZMQ stream (same source as GUI) + Fallback: EPICS area_detector + + Args: + gray: If True, return grayscale image + + Returns: + numpy array with the image + """ + # Try ZMQ source first if configured + if self.__use_zmq_camera and self.__zmq_camera is not None: + logger.debug("we are using zmq camera for mlbox and screenshots") + image = self.__zmq_camera.get_image(gray=gray) + if image is not None: + return image + logger.warning("ZMQ camera unavailable, falling back to EPICS area_detector") + + # Fallback to EPICS area_detector return self.__sample_cam.get_image(gray=gray) def samcam_auto(self, state: AutoEnum): @@ -269,6 +297,22 @@ class BeamlineDevices: """ return int(self.__sample_cam.uid.get()) + @property + def zmq_camera_enabled(self) -> bool: + """Check if ZMQ camera source is enabled.""" + return self.__use_zmq_camera + + @zmq_camera_enabled.setter + def zmq_camera_enabled(self, enabled: bool): + """Enable or disable ZMQ camera source (fallback to EPICS when disabled).""" + self.__use_zmq_camera = enabled and self.__zmq_camera is not None + logger.info(f"ZMQ camera source {'enabled' if self.__use_zmq_camera else 'disabled'}") + + @property + def zmq_camera_connected(self) -> bool: + """Check if ZMQ camera is currently connected.""" + return self.__zmq_camera is not None and self.__zmq_camera.is_connected() + # Detector Z @property def dtz(self) -> float: diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 0705b15e..19920037 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -1816,6 +1816,19 @@ async def send_screenshot_db( daq.send_screenshot_db(filename=filename, message=message) return "OK" +@app.get("/camera/source") +async def get_camera_source(token: str = Depends(oauth2_scheme)): + """Get the current camera image source.""" + auth.check_jwt_ro(cfg, auth.parse_token(token)) + return {"source": daq._AareDAQ__devs.samcam_source} + + +@app.post("/camera/source") +async def set_camera_source(use_zmq: bool, token: str = Depends(oauth2_scheme)): + """Set the camera image source preference.""" + auth.check_jwt_rw(cfg, auth.parse_token(token)) + daq._AareDAQ__devs.set_camera_source(use_zmq) + return {"source": daq._AareDAQ__devs.samcam_source} LOGGING_CONFIG = { "version": 1, diff --git a/src/aare/devices/zmq_client.py b/src/aare/devices/zmq_client.py new file mode 100644 index 00000000..c4a7bf6b --- /dev/null +++ b/src/aare/devices/zmq_client.py @@ -0,0 +1,186 @@ +""" +ZMQ-based camera image source. +Provides on-demand image retrieval via ZMQ REQ/REP pattern, +with fallback to area_detector if ZMQ is unavailable. +""" + +import json +import time +from typing import Optional + +import cv2 +import numpy as np +import zmq + +from aare.common.beamline import MXBeamline +from aare.common.logger_config import setup_logger + +logger = setup_logger("aareDAQ") + + +class ZMQCameraClient: + """ + ZMQ camera client that requests the latest image from a ZMQ stream. + + Uses a SUB socket with a short timeout to grab the most recent frame, + rather than maintaining a continuous subscription. This is suitable + for on-demand image retrieval in the DAQ server. + """ + + def __init__(self, beamline: MXBeamline, timeout_ms: int = 2000): + """ + Initialize the ZMQ camera client. + + Args: + zmq_url: ZMQ endpoint URL (e.g., "tcp://x10sa-spark-01:9091") + timeout_ms: Receive timeout in milliseconds + """ + if beamline == MXBeamline.X06DA: + self.__simulated = False + self.__zmq_url = "tcp://x06da-pserv-01:9089" + elif beamline == MXBeamline.X10SA: + self.__simulated = False + self.__zmq_url = "tcp://x10sa-spark-01:9091" + elif beamline == MXBeamline.SIMULATED: + self.__simulated = True + self.__zmq_url = None + else: + raise Exception("unknown beamline") + + self.__timeout_ms = timeout_ms + self.__context: Optional[zmq.Context] = None + self.__socket: Optional[zmq.Socket] = None + self.__last_image: Optional[np.ndarray] = None + self.__last_fetch_time: float = 0.0 + self.__connected = False + + def _ensure_connected(self) -> bool: + """Ensure the ZMQ socket is connected. Returns True if successful.""" + if self.__connected and self.__socket is not None: + return True + + if self.__simulated or self.__zmq_url is None: + return False + + try: + if self.__context is None: + self.__context = zmq.Context() + + if self.__socket is not None: + self.__socket.close() + + self.__socket = self.__context.socket(zmq.SUB) + self.__socket.setsockopt(zmq.SUBSCRIBE, b"") + self.__socket.setsockopt(zmq.RCVTIMEO, self.__timeout_ms) + # Discard old messages, keep only the latest + self.__socket.setsockopt(zmq.CONFLATE, 1) + self.__socket.connect(self.__zmq_url) + self.__connected = True + logger.debug(f"ZMQ camera connected to {self.__zmq_url}") + return True + except Exception as e: + logger.warning(f"Failed to connect ZMQ camera to {self.__zmq_url}: {e}") + self.__connected = False + return False + + def get_image(self, gray: bool = False) -> Optional[np.ndarray]: + """ + Fetch the latest image from the ZMQ stream. + + Args: + gray: If True, return grayscale image; otherwise return RGB. + + Returns: + numpy array with the image, or None if unavailable. + """ + if self.__simulated: + return None + + if not self._ensure_connected(): + return None + + try: + r = self.__socket.recv_multipart() + + if len(r) < 2: + logger.debug("ZMQ camera: incomplete message received") + return self.__last_image + + data = r[-1] + header = None + + for part in r[:-1]: + try: + decoded = json.loads(part.decode("utf-8")) + if isinstance(decoded, dict) and "shape" in decoded: + header = decoded + break + except Exception: + continue + + if header and header.get("type") == "uint8": + h, w = header["shape"][:2] + raw = np.frombuffer(data, np.uint8).reshape((h, w)) + + # Convert from Bayer to RGB + rgb = cv2.cvtColor(raw, cv2.COLOR_BAYER_GB2RGB) + + if gray: + gray_img = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + self.__last_image = gray_img + return gray_img + else: + self.__last_image = rgb + return rgb + else: + logger.debug("ZMQ camera: unexpected image format") + return self.__last_image + + except zmq.Again: + logger.debug("ZMQ camera: timeout waiting for frame") + return self.__last_image + except Exception as e: + logger.warning(f"ZMQ camera error: {e}") + self.__connected = False + return self.__last_image + + def is_available(self) -> bool: + """Check if ZMQ camera is available and responding.""" + if self.__simulated: + return False + + if not self._ensure_connected(): + return False + + try: + old_timeout = self.__socket.getsockopt(zmq.RCVTIMEO) + self.__socket.setsockopt(zmq.RCVTIMEO, 500) + try: + self.__socket.recv_multipart() + self.__socket.setsockopt(zmq.RCVTIMEO, old_timeout) + return True + except zmq.Again: + self.__socket.setsockopt(zmq.RCVTIMEO, old_timeout) + return False + except Exception: + return False + + def close(self): + """Close the ZMQ connection.""" + if self.__socket: + try: + self.__socket.close() + except Exception: + pass + self.__socket = None + if self.__context: + try: + self.__context.term() + except Exception: + pass + self.__context = None + self.__connected = False + + @property + def url(self) -> str: + return self.__zmq_url if self.__zmq_url else "simulated" \ No newline at end of file