From f724e7e69ee986e01b82bd074dd1071a9cecc504 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 23 Sep 2025 09:41:47 +0200 Subject: [PATCH] Adding prediction zmq and drawing in camera image --- gui/src/aaregui/gui.py | 19 ++++++- gui/src/aaregui/main_window.py | 16 +++++- gui/src/aaregui/threads/camera_thread.py | 42 ++++++++++++++ gui/src/aaregui/widgets/camera_image.py | 71 ++++++++++++++++++++++++ 4 files changed, 145 insertions(+), 3 deletions(-) diff --git a/gui/src/aaregui/gui.py b/gui/src/aaregui/gui.py index f8779aef..cce33c5b 100644 --- a/gui/src/aaregui/gui.py +++ b/gui/src/aaregui/gui.py @@ -25,15 +25,19 @@ if __name__ == "__main__": case MXBeamline.X06DA: default_url = "http://x06da-queue-01.psi.ch:5210" default_zmq_addr = "tcp://x06da-pserv-01:9089" #129.129.110.12:9089 + default_pred_zmq_addr = "tcp://mx-ml:9091" case MXBeamline.X10SA: default_url = "http://x10aa-queue-01.psi.ch:5210" default_zmq_addr = "" + default_pred_zmq_addr = "" case MXBeamline.X06SA: default_url = "http://x06sa-queue-01.psi.ch:5210" default_zmq_addr = "" + default_pred_zmq_addr = "" case _: default_url = "" default_zmq_addr = "" + default_pred_zmq_addr = "" # Add custom options as needed urlOption = QCommandLineOption(["u", "aaredaq-url"], @@ -52,6 +56,14 @@ if __name__ == "__main__": default_zmq_addr) parser.addOption(cameraZeroMQ) + predZmqOption = QCommandLineOption( + ["p", "pred-zmq"], + "Prediction ZeroMQ URL (PUB) to subscribe to (e.g. tcp://mx-ml:9091)", + "pred-zmq", + default_pred_zmq_addr + ) + parser.addOption(predZmqOption) + parser.process(app) base_url = parser.value(urlOption) @@ -62,6 +74,10 @@ if __name__ == "__main__": if zmq_addr == "": zmq_addr = None + pred_zmq_addr = parser.value(predZmqOption) + if pred_zmq_addr == "": + pred_zmq_addr = None + if parser.isSet(defaultImage): default_image = parser.value(defaultImage) else: @@ -73,7 +89,8 @@ if __name__ == "__main__": win = MainWindow(base_url=base_url, token=dialog.token, default_image=default_image, - zmq_addr=zmq_addr) + zmq_addr=zmq_addr, + pred_zmq_addr=pred_zmq_addr) win.show() sys.exit(app.exec()) else: diff --git a/gui/src/aaregui/main_window.py b/gui/src/aaregui/main_window.py index eca43739..f94384af 100644 --- a/gui/src/aaregui/main_window.py +++ b/gui/src/aaregui/main_window.py @@ -25,7 +25,7 @@ from aaregui.scan_logic.raster_grid_manager import RasterGridManager from aaregui.scan_logic.rotation_scan_manager import RotationScanManager from aaregui.scan_logic.sample_mount_logic import SampleMountLogic from aaregui.threads.axis_video_thread import VideoThread -from aaregui.threads.camera_thread import SampleCameraThread +from aaregui.threads.camera_thread import SampleCameraThread, PredictionSubscriber from aaregui.threads.daq_worker import DAQWorker from aaregui.threads.jfjoch_viewer import JFJochDBusClient from aaregui.widgets.camera_image import SampleCameraImageLabel @@ -40,7 +40,8 @@ class MainWindow(QMainWindow): def __init__(self, base_url: str | None, token: str, default_image: str | None, - zmq_addr: str | None): + zmq_addr: str | None, + pred_zmq_addr: str | None): super().__init__() self.__base_url = base_url self.__token = token @@ -188,6 +189,15 @@ class MainWindow(QMainWindow): else: self.camera_thread = None + # Prediction subscriber thread + if pred_zmq_addr is not None: + self.prediction_thread = PredictionSubscriber(pred_zmq_url=pred_zmq_addr) + self.prediction_thread.prediction.connect(self.sample_camera.update_detections) + self.prediction_thread.start() + else: + self.prediction_thread = None + + QApplication.instance().aboutToQuit.connect(self.cleanup) # @@ -344,5 +354,7 @@ class MainWindow(QMainWindow): def cleanup(self): if self.camera_thread is not None: self.camera_thread.stop() + if self.prediction_thread is not None: + self.prediction_thread.stop() if self.beamline_camera_thread is not None: self.beamline_camera_thread.stop() diff --git a/gui/src/aaregui/threads/camera_thread.py b/gui/src/aaregui/threads/camera_thread.py index 4633fa75..d4bfa6c1 100644 --- a/gui/src/aaregui/threads/camera_thread.py +++ b/gui/src/aaregui/threads/camera_thread.py @@ -52,3 +52,45 @@ class SampleCameraThread(QThread): self.quit() print("Camera thread stopped. Exiting...") 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/gui/src/aaregui/widgets/camera_image.py b/gui/src/aaregui/widgets/camera_image.py index 599b6324..1d03ef04 100644 --- a/gui/src/aaregui/widgets/camera_image.py +++ b/gui/src/aaregui/widgets/camera_image.py @@ -108,6 +108,19 @@ class SampleCameraImageLabel(QGraphicsView): self.left_click_hold_threshold = 200 self.right_click_hold_threshold = 200 + self.__detections = [] # list of dicts from publisher + self.__det_shape = None # shape from payload [h,w] so we can scale + + @Slot(dict) + def update_detections(self, payload: dict): + # payload: { 'time', 'frame_id', 'shape':[h,w], 'boxes':[{'x1',...,'label','conf'}] } + try: + self.__det_shape = payload.get('shape', None) + self.__detections = payload.get('boxes', []) or [] + except Exception: + self.__detections = [] + self.update() # trigger redraw + def __draw_busy_overlay(self, painter: QPainter, rect): """Draw 'BEAMLINE BUSY' text overlay when DAQ is in busy state.""" if not self.__is_daq_busy: @@ -165,6 +178,7 @@ class SampleCameraImageLabel(QGraphicsView): self.__raster_mgr.draw_grid(painter, self.__raster_alpha) self.__draw_helical(painter) self.__draw_busy_overlay(painter, rect) + self.__draw_detections(painter, rect) def resizeEvent(self, event): super().resizeEvent(event) @@ -402,6 +416,63 @@ class SampleCameraImageLabel(QGraphicsView): ) self.smargon.emit(smargon_coord) + def __draw_detections(self, painter: QPainter): + if not getattr(self, "_SampleCameraImageLabel__detections", None): + return + if self.pixmap_item is None: + return + pix = self.pixmap_item.pixmap() + if pix.isNull(): + return + + # size on disk / image from payload + det_shape = self.__det_shape # [h, w] + if det_shape is None: + return + img_h, img_w = det_shape[0], det_shape[1] + # displayed pixmap size (in pixels) + disp_w = pix.width() + disp_h = pix.height() + # scale factors from model/image -> displayed pixmap coordinates + sx = disp_w / float(img_w) + sy = disp_h / float(img_h) + + # color mapping requested + color_map = { + 'pin': QColor('red'), + 'loop_all': QColor('green'), + 'loop_face': QColor('yellow'), + 'crystal': QColor('blue'), + } + + pen = QPen() + pen.setWidth(2) + painter.setBrush(Qt.BrushStyle.NoBrush) + + for det in self.__detections: + try: + x1 = det['x1'] * sx + y1 = det['y1'] * sy + x2 = det['x2'] * sx + y2 = det['y2'] * sy + label = str(det.get('label', '')).lower() + conf = det.get('conf', 0.0) + except Exception: + continue + + color = color_map.get(label, QColor('magenta')) + pen.setColor(color) + painter.setPen(pen) + rect = QRect(int(x1), int(y1), int(max(1, x2 - x1)), int(max(1, y2 - y1))) + painter.drawRect(rect) + # label text background for readability + painter.setPen(QPen(QColor(0, 0, 0), 0)) + painter.setBrush(color) + # draw small background rectangle then label text + painter.drawRect(QRect(int(x1), int(y1 - 16), int(8 + 7 * len(label)), 16)) + painter.setPen(QPen(QColor(0, 0, 0))) + painter.drawText(QPoint(int(x1) + 2, int(y1 - 4)), f"{label} {conf:.2f}") + def __draw_ml_bounding_box(self, painter: QPainter): if self.__bounding_box is None: return