From 731efccc2bacdafb72c88fe0f2d3765efdf08e85 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 4 Jul 2025 10:59:42 +0200 Subject: [PATCH] DAQ: added ml_bounding_box to loop centering, creates grid scan grid using ml. Fixed bug where filename wouldn't iterate for copied grid scans. --- common/src/aaredaqlib/models.py | 7 +++ daq/src/aaredaq/daq.py | 44 ++++++++++++++----- daq/src/aaredaq/server.py | 18 +++++--- gui/src/aaregui/main_window.py | 3 ++ .../aaregui/panels/loop_centering_panel.py | 7 ++- .../aaregui/scan_logic/raster_grid_manager.py | 20 ++++++++- gui/src/aaregui/threads/daq_worker.py | 23 ++++++++++ gui/src/aaregui/widgets/camera_image.py | 20 +++++++++ 8 files changed, 124 insertions(+), 18 deletions(-) diff --git a/common/src/aaredaqlib/models.py b/common/src/aaredaqlib/models.py index 71f68053..5fb55bfb 100644 --- a/common/src/aaredaqlib/models.py +++ b/common/src/aaredaqlib/models.py @@ -94,6 +94,12 @@ class BeamMarkCoeffModel(BaseModel): y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2], ) +class BoundingBoxModel(BaseModel): + top_x: float + top_y: float + bottom_x: float + bottom_y: float + class LoopCenteringZoomModelElem(BaseModel): zoom_value: float @@ -177,6 +183,7 @@ class DAQStatusModel(BaseModel): busy: bool sample: SampleShortInfo | None = None session: SessionStatus + box: BoundingBoxModel | None = None class BeamlineSettingsModel(BaseModel): dtz_max: float | None = 1600.0 diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index 07e2c9dd..ead4e181 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -21,7 +21,7 @@ from aaredaqlib.models import ( SampleShortInfo, PuckLoadedInfo, SampleShortInfoList, - DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, ) + DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, BoundingBoxModel, ) from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid from aaredaqlib.rotation_scan import RotationScanRequest, CompletedRotationScan from aaredaqlib.sample_geometry import SampleGeometryModel @@ -50,7 +50,8 @@ class AareDAQ: self.__jfjoch = JFJochWrapper(bl) self.__bl = bl.value.upper() self.__aare = AareWrapper(bl) - + self.__saved_box = None + @property def state(self) -> BeamlineStateEnum: return self.__cfg.state @@ -80,6 +81,7 @@ class AareDAQ: @omega.setter def omega(self, val: float): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) + self.__saved_box = None if -720 < val < 720: self.__devs.aerotech.move(val, wait=True, speed=180.0, direct=True) self.__cfg.state_busy = False @@ -93,6 +95,7 @@ class AareDAQ: @zoom.setter def zoom(self, val: float): + self.__saved_box = None self.__devs.zoom = val @property @@ -153,6 +156,7 @@ class AareDAQ: new_meas_pos = self.__cfg.abr_meas_pos + c self.__cfg.abr_meas_pos = new_meas_pos self.__devs.abr_pos = new_meas_pos + self.__saved_box = None self.__cfg.state_busy = False except: self.__cfg.state_busy = False @@ -193,6 +197,7 @@ class AareDAQ: def __mount(self, target: SampleShortInfo | None): self.__set_state(BeamlineStateEnum.RobotSampleExchange) + self.__saved_box = None self.__devs.smargon.move_home(wait=True) self.__devs.abr_pos = ABR_POS_MOUNT time.sleep(0.5) @@ -433,6 +438,7 @@ class AareDAQ: def smargon(self, sc: SmargonCoordinate): self.__cfg.set_busy(BeamlineStateEnum.SampleAlignment) try: + self.__saved_box = None self.__devs.smargon.target = sc self.__devs.smargon.wait() self.__cfg.state_busy = False @@ -495,7 +501,7 @@ class AareDAQ: def listen_changes(self) -> redis.client.PubSub: return self.__cfg.listen_changes() - def __ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None, move: bool = False) -> None | Tuple[float, float, float, float]: + def __ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None, move: bool = False) -> BoundingBoxModel | None: time.sleep(0.2) # Just to be sure image is stable #curr_image = self.camera_image #box = self.__mlbox.predict(curr_image) @@ -515,16 +521,33 @@ class AareDAQ: coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1)) self.__devs.smargon.target = SmargonCoordinate(sh_mm=coord) self.__devs.smargon.wait(60) - return box + return BoundingBoxModel(bottom_x=x1, bottom_y=y1, top_x=x2, top_y=y2) - def ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None, move: bool = False) -> None | Tuple[float, float, float, float]: + def ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None, move: bool = False) -> RasterGridRequest | None: try: self.__cfg.try_set_busy(timeout=360) - box = self.__ml_bounding_box(sample_id, filename, move) - if not box: - raise LoopCenteringFailed + ml_box = self.__ml_bounding_box(sample_id, filename, move) + r = None + if ml_box is not None: + geom = self.sample_geometry + + start_coord = geom.picture_to_smargon(Coordinate(x=ml_box.bottom_x, y=ml_box.bottom_y)) + grid_size = Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.8) + n_x = abs(ceil((ml_box.top_x - ml_box.bottom_x) * geom.pixel_in_mm / grid_size.x)) + n_y = abs(ceil((ml_box.top_y - ml_box.bottom_y) * geom.pixel_in_mm / grid_size.y)) + + r = RasterGridRequest( + exp_time_s=0.05, + smargon= SmargonCoordinate(chi_deg = geom.smargon.chi_deg, + phi_deg= geom.smargon.phi_deg, + sh_mm=start_coord), + n_x=n_x, + n_y=n_y, + grid_size_mm= grid_size, + omega_deg=geom.omega_deg + ) self.__cfg.state_busy = False - return box + return r except Exception: self.__cfg.state_busy = False raise @@ -836,7 +859,8 @@ class AareDAQ: bl=self.beamline_status, sample=self.sample, session=SessionStatus(), - diffraction=self.diffraction_geometry + diffraction=self.diffraction_geometry, + box=self.__saved_box ) def cancel(self): diff --git a/daq/src/aaredaq/server.py b/daq/src/aaredaq/server.py index cf83b7c5..661fdc78 100644 --- a/daq/src/aaredaq/server.py +++ b/daq/src/aaredaq/server.py @@ -7,7 +7,8 @@ import urllib3 import uvicorn from aaredaqlib.coordinate import SmargonCoordinate, Coordinate from aaredaqlib.models import SampleShortInfo, DAQStatusModel, BeamlineStateEnum, BeamlineSettingsModel, \ - SampleShortInfoList, SessionStatus, SampleCameraSettings, AutofocusSettings, TokenData + SampleShortInfoList, SessionStatus, SampleCameraSettings, AutofocusSettings, TokenData, \ + CryojetSettingsModel from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid from aaredaqlib.rotation_scan import RotationScanRequest, CompletedRotationScan from aaredaqlib.sample_geometry import SampleGeometryModel @@ -330,11 +331,9 @@ async def alc_center_loop(token: str = Depends(oauth2_scheme)) -> str: @app.post("/alc/ml_bounding_box") -async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme), move: bool = False, filename: str = "") -> Tuple[ - float, float, float, float] | None: +async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme)) -> RasterGridRequest | None: auth.check_jwt_rw(cfg, auth.parse_token(token)) - return daq.ml_bounding_box(move=move) - + return daq.ml_bounding_box() # Access management @app.get("/access/pgroup") @@ -387,6 +386,15 @@ async def put_settings(s: BeamlineSettingsModel, token: str = Depends(oauth2_sch auth.check_jwt_staff(auth.parse_token(token)) cfg.settings = s +@app.get("/beamline/cryo_settings") +async def get_cryo_settings(token: str = Depends(oauth2_scheme)) -> CryojetSettingsModel: + auth.check_jwt_staff(auth.parse_token(token)) + return cfg.cryojet_settings + +@app.put("/beamline/cryo_settings") +async def put_cryo_settings(s: CryojetSettingsModel, token: str = Depends(oauth2_scheme)): + auth.check_jwt_staff(auth.parse_token(token)) + cfg.cryojet_settings = s LOGGING_CONFIG = { "version": 1, diff --git a/gui/src/aaregui/main_window.py b/gui/src/aaregui/main_window.py index d44c41be..28c8cb63 100644 --- a/gui/src/aaregui/main_window.py +++ b/gui/src/aaregui/main_window.py @@ -143,6 +143,9 @@ class MainWindow(QMainWindow): self.beamline.samcam.changed.connect(self.daq.samcam_settings) self.beamline.loopctr.background.clicked.connect(self.daq.alc_background) self.beamline.loopctr.find_tip.clicked.connect(self.daq.center_loop) + self.beamline.loopctr.bounding_box.clicked.connect(self.daq.ml_bounding_box) + self.daq.raster_generated_by_ml.connect(self.raster.update_active_grid_request) + self.camera_image.smargon.connect(self.daq.move_smargon) self.beamline.smargon_panel.smargon.connect(self.daq.move_smargon) self.camera_image.samcam_updated.connect(self.daq.samcam_settings) diff --git a/gui/src/aaregui/panels/loop_centering_panel.py b/gui/src/aaregui/panels/loop_centering_panel.py index c113b53a..18199693 100644 --- a/gui/src/aaregui/panels/loop_centering_panel.py +++ b/gui/src/aaregui/panels/loop_centering_panel.py @@ -8,11 +8,14 @@ class LoopCenteringPanel(QWidget): super().__init__(parent) grid_layout = QGridLayout(self) - grid_layout.addWidget(TitleLabel("Loop centering", self), 0, 0, 1, 2) + grid_layout.addWidget(TitleLabel("Loop centering", self), 0, 0, 1, 3) grid_layout.setColumnStretch(0, 1) grid_layout.setColumnStretch(1, 1) self.find_tip = QPushButton("Center", parent=self) grid_layout.addWidget(self.find_tip, 1, 0) - self.background = QPushButton("Background", parent=self) + self.background = QPushButton("Bkg", parent=self) grid_layout.addWidget(self.background, 1, 1) + + self.bounding_box = QPushButton("Box", parent=self) + grid_layout.addWidget(self.bounding_box, 1, 2) diff --git a/gui/src/aaregui/scan_logic/raster_grid_manager.py b/gui/src/aaregui/scan_logic/raster_grid_manager.py index e824f430..beb3e867 100644 --- a/gui/src/aaregui/scan_logic/raster_grid_manager.py +++ b/gui/src/aaregui/scan_logic/raster_grid_manager.py @@ -123,6 +123,18 @@ class RasterGridManager(QObject): return True return False + @Slot(RasterGridRequest) + def update_active_grid_request(self, grid: RasterGridRequest): + self.__active_grid.smargon = grid.smargon + self.__active_grid.omega_deg = grid.omega_deg + self.__active_grid.grid_size_mm = grid.grid_size_mm + self.__active_grid.n_x = grid.n_x + self.__active_grid.n_y = grid.n_y + self.grid_scan_size_changed.emit(self.__active_grid.n_x, + self.__active_grid.n_y, + self.__active_grid.grid_size_mm.x, + self.__active_grid.grid_size_mm.y) + @Slot(DAQStatusModel) def update_daq_status(self, s: DAQStatusModel): if s.sample is not None and s.sample.db_id != self.__sample_id: @@ -401,7 +413,13 @@ class RasterGridManager(QObject): @Slot(int) def completed_grid_redo(self, row: int): if 0 <= row < len(self.__completed_grids): - self.__active_grid = copy.deepcopy(self.__completed_grids[row].request) + copy_of_active_grid = copy.deepcopy(self.__completed_grids[row].request) + self.__active_grid.n_x = copy_of_active_grid.n_x + self.__active_grid.n_y = copy_of_active_grid.n_y + self.__active_grid.grid_size_mm = copy_of_active_grid.grid_size_mm + self.__active_grid.smargon = copy_of_active_grid.smargon + self.__active_grid.omega_deg = copy_of_active_grid.omega_deg + self.__active_grid.visible = True self.__completed_grids[row].request.visible = False diff --git a/gui/src/aaregui/threads/daq_worker.py b/gui/src/aaregui/threads/daq_worker.py index fa3bfacc..576098ef 100644 --- a/gui/src/aaregui/threads/daq_worker.py +++ b/gui/src/aaregui/threads/daq_worker.py @@ -23,6 +23,7 @@ class DAQWorker(QObject): run_number_incremented = Signal() raster_scan_completed = Signal(CompletedRasterGrid) standard_scan_completed = Signal(CompletedRotationScan) + raster_generated_by_ml = Signal(RasterGridRequest) def __init__(self, base_url: str | None, token: str, parent=None): super().__init__(parent) @@ -342,3 +343,25 @@ class DAQWorker(QObject): @Slot() def cancel(self): self.generic_post("scan/cancel") + + def handle_ml_box_response(self, reply): + try: + response_data = self.handle_response(reply) + if response_data != "": + parsed_response = RasterGridRequest.model_validate_json(response_data) + self.raster_generated_by_ml.emit(parsed_response) + except Exception as e: + print(e) + self.http_error.emit(str(e)) + + @Slot() + def ml_bounding_box(self): + if self.__base_url is None: + print(f"POST /alc/ml_bounding_box") + return + + request = QNetworkRequest(QUrl(f"{self.__base_url}/alc/ml_bounding_box")) + request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8")) + reply = self.__net_manager.post(request, QByteArray(b"")) + reply.finished.connect(lambda: self.handle_ml_box_response(reply)) + diff --git a/gui/src/aaregui/widgets/camera_image.py b/gui/src/aaregui/widgets/camera_image.py index a2a06d1d..abb3b9cf 100644 --- a/gui/src/aaregui/widgets/camera_image.py +++ b/gui/src/aaregui/widgets/camera_image.py @@ -72,6 +72,7 @@ class SampleCameraImageLabel(QGraphicsView): self.__helical_start = SmargonCoordinate() self.__helical_end = SmargonCoordinate() self.__raster_alpha = 127 + self.__bounding_box = None self.start_point = None # Starting point of the rectangle self.end_point = None # Ending point of the rectangle @@ -106,6 +107,7 @@ class SampleCameraImageLabel(QGraphicsView): self.right_click_hold_threshold = 200 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) @@ -309,6 +311,8 @@ class SampleCameraImageLabel(QGraphicsView): def update_daq_status(self, s: DAQStatusModel): self.__geom = s.geom self.__sam_cam = s.bl.sample_camera + self.__bounding_box = s.box + if s.state == BeamlineStateEnum.BeamLocation: self.__state = SampleCameraImageState.BEAM_MARKING elif self.__state == SampleCameraImageState.BEAM_MARKING: @@ -332,6 +336,22 @@ class SampleCameraImageLabel(QGraphicsView): ) self.smargon.emit(smargon_coord) + 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