mx/AareDAQ
0
1
Fork 0
Code Issues 49 Pull Requests 1 Actions 1 Packages Projects Releases 82 Wiki Activity

Adding prediction zmq and drawing in camera image

Browse Source
This commit is contained in:
GotthardG
2025-09-23 09:41:47 +02:00
parent d90d847506
commit f724e7e69e
4 changed files with 145 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gui/src/aaregui/gui.py
+18 -1
View File
@@ -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:
gui/src/aaregui/main_window.py
+14 -2
View File
@@ -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()
gui/src/aaregui/threads/camera_thread.py
+42
View File
@@ -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()
gui/src/aaregui/widgets/camera_image.py
+71
View File
@@ -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