175 lines
6.7 KiB
Python
175 lines
6.7 KiB
Python
import json
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import zmq
|
|
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
|
|
|
|
|
|
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:
|
|
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:
|
|
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)
|