AareDAQ/src/aare/gui/threads/camera_thread.py

import json
import time

import cv2
import numpy as np
import zmq
from aarecommon.config.logger import setup_logger
from aarecommon.math.autofocus import focus_measure_edges
from aarecommon.models.models import DAQStatusModel
from PySide6.QtCore import QThread, Signal, Slot
from PySide6.QtGui import QImage, QPixmap

from aare.gui.constants import LOGGER_NAME

logger = setup_logger(LOGGER_NAME)


class SampleCameraThread(QThread):
    # Define a signal to communicate messages from the thread to the main GUI
    camera_image = Signal(QPixmap)
    focus_measure = Signal(float)  # Emits Laplacian variance (higher = sharper)
    fps_measure = Signal(float)
    camera_availability_changed = Signal(bool)
    camera_error = Signal(str)

    def __init__(self, zmq_url: str, parent=None):
        super().__init__(parent)
        context = zmq.Context()
        self._socket = context.socket(zmq.SUB)
        self._socket.setsockopt(zmq.SUBSCRIBE, b"")
        self._socket.setsockopt(zmq.RCVTIMEO, 500)
        self._socket.connect(zmq_url)

        self.running = True
        self._measure_focus = False
        self._focus_mask = None
        self._beam_x = 0
        self._beam_y = 0
        self._last_beam_pos = None
        self._radius = 40

        self._fps_window_start = time.perf_counter()
        self._fps_frame_count = 0
        self._fps_emit_period_s = 0.5
        self._last_frame_time = None
        self._no_frame_timeout_s = 5.0

        self._camera_available = False
        self._last_camera_error: str | None = None

    def _set_camera_available(self, available: bool, error: str | None = None) -> None:
        if available != self._camera_available:
            self._camera_available = available
            self.camera_availability_changed.emit(available)

        if error is not None and error != self._last_camera_error:
            self._last_camera_error = error
            self.camera_error.emit(error)

        if available:
            self._last_camera_error = None

    @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):
        """Enable or disable focus measurement."""
        self._measure_focus = enabled

    def run(self):
        while self.running:
            try:
                r = self._socket.recv_multipart()

                now = time.perf_counter()
                self._last_frame_time = now
                self._fps_frame_count += 1

                elapsed = now - self._fps_window_start
                if elapsed >= self._fps_emit_period_s:
                    fps = self._fps_frame_count / elapsed if elapsed > 0 else 0.0
                    self.fps_measure.emit(float(fps))
                    self._fps_window_start = now
                    self._fps_frame_count = 0

                if len(r) < 2:
                    continue

                data = r[-1]
                header = None

                for part in r[:-1]:
                    try:
                        decoded = json.loads(part.decode("utf-8"))
                        if isinstance(decoded, dict):
                            header = decoded
                            break
                    except Exception:
                        logger.debug(
                            "Skipping an unparsable camera ZMQ message part", exc_info=True
                        )
                        continue

                if header:
                    encoding = str(header.get("encoding", "")).lower()

                    if encoding == "jpeg":
                        encoded = np.frombuffer(data, dtype=np.uint8)
                        bgr = cv2.imdecode(encoded, cv2.IMREAD_COLOR)
                        if bgr is None:
                            self._set_camera_available(
                                False, "Sample camera feed unavailable: failed to decode JPEG frame"
                            )
                            continue
                        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
                    elif "shape" in header:
                        h, w = header["shape"][:2]
                        raw = np.frombuffer(data, np.uint8).reshape((h, w))
                        rgb = cv2.cvtColor(raw, cv2.COLOR_BAYER_GB2RGB)
                    else:
                        self._set_camera_available(
                            False,
                            f"Sample camera feed unavailable: unsupported frame header {header}",
                        )
                        continue

                    self._set_camera_available(True)

                    if self._measure_focus:
                        gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
                        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

                        sharpness = focus_measure_edges(gray, self._focus_mask)
                        self.focus_measure.emit(sharpness)

                    qimage = QImage(
                        rgb.data, rgb.shape[1], rgb.shape[0], QImage.Format.Format_RGB888
                    ).copy()
                    self.camera_image.emit(QPixmap.fromImage(qimage))
                else:
                    self._set_camera_available(
                        False, "Sample camera feed unavailable: no frame header in zmq stream"
                    )
            except zmq.Again:  # Timeout occurred
                now = time.perf_counter()
                elapsed = now - self._fps_window_start
                if elapsed >= self._fps_emit_period_s:
                    no_frames_long = (
                        self._last_frame_time is None
                        or (now - self._last_frame_time) >= self._no_frame_timeout_s
                    )
                    self.fps_measure.emit(float("nan") if no_frames_long else 0.0)
                    self._fps_window_start = now
                    self._fps_frame_count = 0

                    if no_frames_long:
                        self._set_camera_available(False, "Sample camera feed unavailable")
                continue  # Check self.running again
            except Exception as e:
                logger.warning("Sample camera feed unavailable", exc_info=True)
                self._set_camera_available(False, f"Sample camera feed unavailable: {e}")
                self.running = False

    def stop(self):
        self.running = False

        if self._socket:
            self._socket.close()

        if self.isRunning():
            self.quit()
            self.wait(2000)