DAQ: zmq_client and functionality added, please test!

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:

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,

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"