Files
AareDAQ/gui/src/aaregui/widgets/camera_image.py
T

674 lines
25 KiB
Python

import math
from enum import Enum
from PySide6.QtCore import Qt, QTimer, QRect, QPoint, Signal, Slot, QPointF
from PySide6.QtGui import (
QPixmap,
QPainter,
QPen,
QColor,
QWheelEvent,
QTransform,
QCursor,
QLinearGradient, QFont, QFontMetrics,
)
from PySide6.QtWidgets import (
QMenu,
QToolTip,
QGraphicsView,
QGraphicsScene,
QGraphicsPixmapItem,
QFileDialog,
QFrame,
)
from aaredaqlib.models import DAQStatusModel, AutofocusSettings, SampleCameraSettings, BeamlineStateEnum, \
SessionsStateEnum
from aaregui.models.bookmark import SmargonBookmarkList
from aaredaqlib.coordinate import Coordinate, SmargonCoordinate
from aaredaqlib.sample_geometry import SampleGeometryModel
from aaregui.scan_logic.raster_grid_manager import RasterGridManager
from aaredaqlib.logger_config import setup_logger
logger = setup_logger(__name__)
class SampleCameraImageState(Enum):
IDLE = 0
DRAWING_RASTER_GRID = 1
MOVING_RASTER_GRID = 2
RESIZE_RASTER_GRID = 3
BEAM_MARKING = 4
class SampleCameraImageLabel(QGraphicsView):
smargon = Signal(SmargonCoordinate)
evaluate_grid = Signal()
clear_grid = Signal()
clear_evaluated_grids = Signal()
load_image = Signal(QPointF)
zoom_change = Signal(int)
set_omega = Signal(float)
set_helical_start = Signal(SmargonCoordinate)
set_helical_end = Signal(SmargonCoordinate)
update_beam_mark = Signal(float, float)
autofocus = Signal(AutofocusSettings)
samcam_updated = Signal(SampleCameraSettings)
switch_raster_grid = Signal()
def __init__(
self,
geom: SampleGeometryModel,
raster: RasterGridManager,
default_image: str | None,
parent=None,
):
super().__init__(parent)
self.__raster_mgr = raster
self.__state = SampleCameraImageState.IDLE
self.__session_state: SessionsStateEnum | None = None
self.__sam_cam = SampleCameraSettings(exposure=0.1, gain=100.0)
self.__is_daq_busy = False
self.__geom = geom
self.__bookmarks: SmargonBookmarkList = SmargonBookmarkList()
self.__autoscale = False
self.__show_coords = False
self.__helical_start = SmargonCoordinate()
self.__helical_end = SmargonCoordinate()
self.__raster_alpha = 127
self.__bounding_box = None
self.__show_detections = True
self.start_point = None # Starting point of the rectangle
self.end_point = None # Ending point of the rectangle
self.scene = QGraphicsScene(self)
self.setScene(self.scene)
self.setMouseTracking(True)
if default_image is not None:
pixmap = QPixmap(default_image)
self.pixmap_item = QGraphicsPixmapItem(pixmap)
self.scene.addItem(self.pixmap_item)
else:
pixmap = QPixmap(2000, 2000)
pixmap.fill(Qt.GlobalColor.white)
self.pixmap_item = QGraphicsPixmapItem(pixmap)
self.scene.addItem(self.pixmap_item)
self.setFrameShape(QFrame.Shape.NoFrame)
self.setRenderHints(QPainter.RenderHint.Antialiasing)
# Create a timer for throttling wheel events
self.wheel_event_timer = QTimer()
self.wheel_event_timer.setSingleShot(True)
self.wheel_event_threshold = 500 # Minimum delay between wheel events in ms
self.last_wheel_event = 0 # Timestamp of the last processed wheel event
self.click_timer = QTimer() # Timer to detect click vs hold
self.click_timer.setSingleShot(True) # The timer runs only once
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'}] }
#print(f"DEBUG: Full payload received: {payload}")
try:
self.__det_shape = payload.get('shape', None)
self.__detections = payload.get('boxes', []) or []
except Exception as e:
logger.debug(f"Exception in update_detections: {e}")
self.__detections = []
self.update()
def __draw_busy_overlay(self, painter: QPainter):
if not self.__is_daq_busy:
return
painter.save()
painter.resetTransform()
font = QFont()
font.setPointSize(24) # Fixed size in points
font.setBold(True)
painter.setFont(font)
painter.setPen(QPen(QColor(255, 255, 255, 200), 2, Qt.PenStyle.SolidLine))
painter.setBrush(QColor(255, 0, 0, 150))
text = "BEAMLINE BUSY"
font_metrics = QFontMetrics(font)
text_rect = font_metrics.boundingRect(text)
position_x = 50 # 50 pixels from left edge
position_y = 50 # 50 pixels from top edge
padding = 20
bg_rect = QRect(
position_x - padding,
position_y - padding,
text_rect.width() + 2 * padding,
text_rect.height() + 2 * padding
)
painter.drawRoundedRect(bg_rect, 10, 10)
painter.setPen(QPen(QColor(255, 255, 255), 2, Qt.PenStyle.SolidLine))
text_pos = QPoint(
position_x,
position_y + font_metrics.ascent()
)
painter.drawText(text_pos, text)
painter.restore()
def __draw_session_overlay(self, painter: QPainter):
if self.__session_state not in (SessionsStateEnum.Vacant, SessionsStateEnum.OwnedByElse):
return
painter.save()
painter.resetTransform()
font = QFont()
font.setPointSize(24)
font.setBold(True)
painter.setFont(font)
if self.__session_state == SessionsStateEnum.Vacant:
bg_color = QColor(255, 215, 0, 180)
text = "Session Vacant"
else:
bg_color = QColor(255, 0, 0, 150)
text = "Session Other"
fm = QFontMetrics(font)
text_rect = fm.boundingRect(text)
padding = 16
vw = self.viewport().width()
vh = self.viewport().height()
bg_w = text_rect.width() + 2 * padding
bg_h = text_rect.height() + 2 * padding
position_x = int((vw - bg_w) / 2)
position_y = int((vh - bg_h) / 2)
bg_rect = QRect(position_x, position_y, bg_w, bg_h)
painter.setPen(QPen(QColor(255, 255, 255, 220)))
painter.setBrush(bg_color)
painter.drawRoundedRect(bg_rect, 10, 10)
painter.setPen(QPen(QColor(255, 255, 255)))
painter.drawText(QPoint(position_x + padding, position_y + padding + fm.ascent()), text)
painter.restore()
def drawForeground(self, painter, rect):
self.__draw_ml_bounding_box(painter)
self.__draw_beam_center(painter)
self.__raster_mgr.draw_grid(painter, self.__raster_alpha)
self.__draw_helical(painter)
self.__draw_busy_overlay(painter)
self.__draw_session_overlay(painter)
self.__draw_detections(painter, rect)
def resizeEvent(self, event):
super().resizeEvent(event)
self.__scaling()
def mousePressEvent(self, event):
self.start_point = self.mapToScene(event.pos())
match self.__state:
case SampleCameraImageState.BEAM_MARKING:
self.click_timer.start(self.right_click_hold_threshold)
case SampleCameraImageState.IDLE:
self.click_timer.start(self.right_click_hold_threshold)
if event.button() == Qt.MouseButton.RightButton:
if self.__raster_mgr.is_part_of_active_grid(self.start_point):
self.__state = SampleCameraImageState.RESIZE_RASTER_GRID
else:
self.__state = SampleCameraImageState.DRAWING_RASTER_GRID
elif event.button() == Qt.MouseButton.LeftButton:
if self.__raster_mgr.is_part_of_active_grid(self.start_point):
self.__state = SampleCameraImageState.MOVING_RASTER_GRID
def _update_grid(self):
match self.__state:
case SampleCameraImageState.DRAWING_RASTER_GRID:
self.switch_raster_grid.emit()
self.__raster_mgr.update_active_grid(self.start_point, self.end_point)
case SampleCameraImageState.MOVING_RASTER_GRID:
self.switch_raster_grid.emit()
self.__raster_mgr.move_active_grid(self.end_point - self.start_point)
self.start_point = self.end_point
case SampleCameraImageState.RESIZE_RASTER_GRID:
self.switch_raster_grid.emit()
self.__raster_mgr.resize_active_grid(self.end_point)
def mouseMoveEvent(self, event):
mouse_pos = self.mapToScene(event.pos())
if self.__state == SampleCameraImageState.IDLE:
if event.modifiers() & Qt.KeyboardModifier.ShiftModifier:
self.load_image.emit(mouse_pos)
else:
txt = self.__raster_mgr.is_part_of_completed_grid(mouse_pos)
if txt != "":
QToolTip.showText(self.mapToGlobal(event.pos()), txt, self)
elif self.__show_coords:
x, y = mouse_pos.x(), mouse_pos.y()
QToolTip.showText(
self.mapToGlobal(event.pos()), f"{x:.0f}, {y:.0f} pxl", self
)
self.end_point = mouse_pos
if not self.click_timer.isActive():
self.end_point = self.mapToScene(event.pos())
self._update_grid()
self.update()
def mouseReleaseEvent(self, event):
if event.button() == Qt.MouseButton.LeftButton:
if self.click_timer.isActive():
self.click_timer.stop()
self.__left_single_click(event)
else:
self._update_grid()
self.update()
if event.button() == Qt.MouseButton.RightButton:
if self.click_timer.isActive():
self.click_timer.stop()
self.__right_click_menu(event)
else:
self._update_grid()
self.update()
if self.__state != SampleCameraImageState.BEAM_MARKING:
self.__state = SampleCameraImageState.IDLE
def __right_click_menu(self, event):
menu = QMenu(self)
scale_action = menu.addAction("Scale to fit")
scale_action.setCheckable(True)
scale_action.setChecked(self.__autoscale)
show_coord_action = menu.addAction("Show coordinates")
show_coord_action.setCheckable(True)
show_coord_action.setChecked(self.__show_coords)
grab_action = menu.addAction("Grab")
grab_with_overlay_action = menu.addAction("Grab with overlay")
autofocus_action = menu.addAction("Auto-focus")
beam_mark_action = menu.addAction("Mark beam center")
delete_action = None
evaluate_action = None
scene_pos = self.mapToScene(event.pos())
menu.addSection("Grid")
if self.__raster_mgr.is_part_of_active_grid(scene_pos):
delete_action = menu.addAction("Delete grid")
evaluate_action = menu.addAction("Evaluate grid")
delete_completed_action = menu.addAction("Delete completed grids")
action = menu.exec_(self.mapToGlobal(event.pos()))
if action is None:
return
elif action == show_coord_action:
self.__show_coords = not self.__show_coords
elif action == scale_action:
self.__autoscale = not self.__autoscale
self.__scaling()
elif action == delete_action:
self.clear_grid.emit()
elif action == evaluate_action:
self.evaluate_grid.emit()
elif action == delete_completed_action:
self.clear_evaluated_grids.emit()
elif action == grab_action:
self.__screenshot_with_dialog(overlay=False)
elif action == grab_with_overlay_action:
self.__screenshot_with_dialog(overlay=True)
elif action == autofocus_action:
c = self.mapToScene(event.pos())
self.autofocus.emit(AutofocusSettings(center_x_pxl=c.x(), center_y_pxl=c.y(),
radius_pxl=100, z_range_um=0.5,
z_steps=25))
elif action == beam_mark_action:
c = self.mapToScene(event.pos())
self.update_beam_mark.emit(c.x(), c.y())
self.update()
def __screenshot_with_dialog(self, overlay: bool):
file_path, _ = QFileDialog.getSaveFileName(
self,
"Save View As",
"",
"JPEG Files (*.jpg; *.jpeg);;All Files (*)",
)
if file_path:
self.__screenshot(file_path=file_path, overlay=overlay)
def __screenshot(self, file_path: str, overlay: bool):
scene_rect = self.scene.sceneRect()
pixmap = QPixmap(scene_rect.size().toSize())
painter = QPainter(pixmap)
self.scene.render(painter)
if overlay:
self.drawForeground(painter, scene_rect)
painter.end()
pixmap.save(file_path, "JPEG")
def __scaling(self):
if not self.__autoscale:
self.setTransform(QTransform())
self.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
self.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
return
if self.pixmap_item is None:
return
self.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
self.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
if self.pixmap_item.boundingRect().width() == 0 or self.pixmap_item.boundingRect().height() == 0:
ratio = 1.0
else:
ratio_w = self.viewport().size().width() / self.pixmap_item.boundingRect().width()
ratio_h = self.viewport().size().height() / self.pixmap_item.boundingRect().height()
ratio = min(ratio_w, ratio_h)
if ratio < 0.1:
ratio = 0.1
if ratio >= 1.0:
# Don't enable scaling when gain in ratio is < 5% (to avoid back-and-forth)
self.setTransform(QTransform())
else:
matrix = QTransform()
matrix.scale(ratio, ratio)
self.setTransform(matrix)
@Slot(QPixmap)
def update_pixmap(self, pixmap: QPixmap):
if self.pixmap_item is not None: # Ensure pixmap_item exists
self.pixmap_item.setPixmap(pixmap) # Update the pixmap in the item
else:
# If no pixmap item exists (rare case), create one
self.pixmap_item = QGraphicsPixmapItem(pixmap)
self.scene.addItem(self.pixmap_item)
self.update() # Request an update to redraw the view
@Slot(DAQStatusModel)
def update_daq_status(self, s: DAQStatusModel):
self.__geom = s.geom
self.__sam_cam = s.bl.sample_camera
self.__bounding_box = s.box
new_session_state = s.session.session if hasattr(s, 'session') else None
if new_session_state != self.__session_state:
self.__session_state = new_session_state
self.update()
if s.busy != self.__is_daq_busy:
self.__is_daq_busy = s.busy
self.update()
if s.state == BeamlineStateEnum.BeamLocation:
self.__state = SampleCameraImageState.BEAM_MARKING
elif self.__state == SampleCameraImageState.BEAM_MARKING:
self.__state = SampleCameraImageState.IDLE
def __left_single_click(self, event):
match self.__state:
case SampleCameraImageState.BEAM_MARKING:
if event.button() == Qt.MouseButton.LeftButton and event.modifiers() & Qt.KeyboardModifier.ShiftModifier:
point = self.mapToScene(event.pos())
self.update_beam_mark.emit(point.x(), point.y())
case SampleCameraImageState.IDLE:
point = self.mapToScene(event.pos())
if not self.scene.sceneRect().contains(point):
logger.error("Point is outside scene bounds")
return
if event.modifiers() & Qt.KeyboardModifier.ShiftModifier:
sc = SmargonCoordinate(sh_mm=self.__geom.smargon_z(point.y()))
self.smargon.emit(sc)
else:
sample_coord = self.__geom.picture_to_sample(Coordinate(x=point.x(), y=point.y()))
smargon_coord = SmargonCoordinate(
sh_mm=self.__geom.beamline_to_smargon(sample_coord)
)
self.smargon.emit(smargon_coord)
def __draw_detections(self, painter: QPainter, rect):
if not self.__show_detections:
return
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'),
}
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 as e:
logger.debug(f"Error in draw detection {det}: {e}")
continue
color = color_map.get(label, QColor('magenta'))
# Draw only the rectangle outline (no fill)
pen = QPen(color, 3)
painter.setPen(pen)
painter.setBrush(Qt.BrushStyle.NoBrush)
detection_rect = QRect(int(x1), int(y1), int(max(1, x2 - x1)), int(max(1, y2 - y1)))
painter.drawRect(detection_rect)
# Draw label text with white text on colored background
painter.setPen(QPen(QColor(255, 255, 255), 1))
painter.setBrush(color)
text_bg_rect = QRect(int(x1), int(y1 - 16), int(8 + 7 * len(label)), 16)
painter.drawRect(text_bg_rect)
painter.setPen(QPen(QColor(255, 255, 255)))
painter.drawText(QPoint(int(x1) + 2, int(y1 - 4)), f"{label} {conf:.2f}")
@Slot(bool)
def set_show_detections(self, show: bool):
"""Slot to enable/disable showing ML detections"""
self.__show_detections = show
self.update() # Trigger redraw
def __draw_ml_bounding_box(self, painter: QPainter):
if self.__bounding_box is None:
return
painter.setPen(QPen(QColor(50,205, 50), 3, Qt.PenStyle.SolidLine))
painter.setBrush(Qt.BrushStyle.NoBrush)
painter.drawRect(
QRect(
int(self.__bounding_box.bottom_x),
int(self.__bounding_box.bottom_y),
int(self.__bounding_box.top_x - self.__bounding_box.bottom_x),
int(self.__bounding_box.top_y - self.__bounding_box.bottom_y),
)
)
def __draw_beam_center(self, painter: QPainter):
beam_size_pxl = self.__geom.beam_size_pxl
if self.__state == SampleCameraImageState.BEAM_MARKING:
painter.setPen(QPen(QColor(102, 51, 153), 3, Qt.PenStyle.SolidLine))
elif self.__is_daq_busy is True:
# Use red color to indicate busy state
painter.setPen(QPen(QColor(255, 0, 0), 3, Qt.PenStyle.SolidLine))
else:
painter.setPen(QPen(QColor(245, 121, 0), 3, Qt.PenStyle.SolidLine))
painter.setBrush(Qt.BrushStyle.NoBrush)
painter.drawRect(
QRect(
int(self.__geom.beam_location_pxl.x - beam_size_pxl.x / 2),
int(self.__geom.beam_location_pxl.y - beam_size_pxl.y / 2),
int(beam_size_pxl.x),
int(beam_size_pxl.y),
)
)
@staticmethod
def __draw_circle(painter, coord: Coordinate, color: QColor, radius=10):
painter.setPen(QPen(Qt.PenStyle.NoPen))
painter.setBrush(color)
painter.drawEllipse(coord.x - radius, coord.y - radius, 2 * radius, 2 * radius)
@staticmethod
def __draw_arrow(painter: QPainter, start: Coordinate, end: Coordinate):
gradient = QLinearGradient()
gradient.setStart(QPointF(start.x, start.y)) # Start of the gradient (green)
gradient.setFinalStop(QPointF(end.x, end.y)) # End of the gradient (red)
gradient.setColorAt(0.0, QColor("green")) # Start color
gradient.setColorAt(1.0, QColor("red")) # End color
pen = QPen()
pen.setBrush(gradient) # Use gradient as the brush for the pen
pen.setWidth(4) # Set pen width
painter.setPen(pen)
painter.drawLine(QPointF(start.x, start.y), QPointF(end.x, end.y))
def __draw_helical(self, painter: QPainter):
if self.__helical_start.sh_mm is not None:
start_pxl = self.__geom.smargon_to_picture(self.__helical_start.sh_mm)
self.__draw_circle(painter, start_pxl, QColor("green"))
else:
start_pxl = None
if self.__helical_end.sh_mm is not None:
end_pxl = self.__geom.smargon_to_picture(self.__helical_end.sh_mm)
self.__draw_circle(painter, end_pxl, QColor("red"))
else:
end_pxl = None
if start_pxl is not None and end_pxl is not None:
self.__draw_arrow(painter, start_pxl, end_pxl)
@Slot(SmargonCoordinate)
def show_helical_start(self, pos: SmargonCoordinate):
self.__helical_start = pos
self.update()
@Slot(SmargonCoordinate)
def show_helical_end(self, pos: SmargonCoordinate):
self.__helical_end = pos
self.update()
@Slot(int)
def raster_alpha(self, value: int):
if 0 <= value <= 255:
self.__raster_alpha = value
def keyPressEvent(self, event):
mouse_global_pos = QCursor.pos()
mouse_view_pos = self.mapFromGlobal(mouse_global_pos)
mouse_scene_pos = self.mapToScene(mouse_view_pos)
smargon_coord = SmargonCoordinate(
sh_mm=self.__geom.picture_to_smargon(
Coordinate(x=mouse_scene_pos.x(), y=mouse_scene_pos.y())
))
if event.key() == Qt.Key.Key_1:
self.set_helical_start.emit(smargon_coord)
elif event.key() == Qt.Key.Key_2:
self.set_helical_end.emit(smargon_coord)
super().keyPressEvent(event)
def wheelEvent(self, event: QWheelEvent):
if not self.wheel_event_timer.isActive():
delta_y = event.angleDelta().y()
match self.__state:
case SampleCameraImageState.BEAM_MARKING:
if event.modifiers() & Qt.KeyboardModifier.AltModifier:
new_exp_time = self.__sam_cam.exposure * (1.0 + math.copysign(0.1, delta_y))
new_settings = SampleCameraSettings(exposure=new_exp_time, gain=self.__sam_cam.gain)
self.samcam_updated.emit(new_settings)
else:
new_exp_time = self.__sam_cam.exposure * (1.0 + math.copysign(0.5, delta_y))
new_settings = SampleCameraSettings(exposure=new_exp_time, gain=self.__sam_cam.gain)
self.samcam_updated.emit(new_settings)
case SampleCameraImageState.IDLE:
if event.modifiers() & Qt.KeyboardModifier.AltModifier:
new_exp_time = self.__sam_cam.exposure * (1.0 + math.copysign(0.1, delta_y))
new_settings = SampleCameraSettings(exposure=new_exp_time, gain=self.__sam_cam.gain)
self.samcam_updated.emit(new_settings)
elif event.modifiers() & Qt.KeyboardModifier.ControlModifier:
new_exp_time = self.__sam_cam.exposure * (1.0 + math.copysign(0.5, delta_y))
new_settings = SampleCameraSettings(exposure=new_exp_time, gain=self.__sam_cam.gain)
self.samcam_updated.emit(new_settings)
elif event.modifiers() & Qt.KeyboardModifier.ShiftModifier:
self.set_omega.emit(self.__geom.omega_deg + math.copysign(10.0, delta_y))
else:
self.set_omega.emit(self.__geom.omega_deg + math.copysign(90.0, delta_y))
# Start the timer to throttle further events
self.wheel_event_timer.start(self.wheel_event_threshold)