diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 8389860a..804ff5a5 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -5,7 +5,8 @@ import traceback from datetime import datetime from math import ceil -from typing import List, Tuple, Optional +from pathlib import Path +from typing import List, Tuple, Optional, Callable import cv2 import numpy as np @@ -35,6 +36,7 @@ from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan from aare.common.sample_geometry import SampleGeometryModel from aare.devices.area_detector import AutoEnum from aare.devices.jfjoch import JFJochWrapper +from aare.devices.mx_lib import clean_filename from aare.common.exception_handler import ( TransformationInvalidException, @@ -60,6 +62,52 @@ class AareDAQ: self.__bl = bl.value.upper() self.__aare = AareWrapper(bl) self.__saved_box = None + self._smargon_trace_path = Path("logs") / "smargon_trace.csv" + self._face_detection_progress_cb: Callable[[dict], None] | None = None + + def set_face_detection_progress_callback(self, cb: Callable[[dict], None] | None) -> None: + self._face_detection_progress_cb = cb + + def _emit_face_detection_progress(self, payload: dict) -> None: + if self._face_detection_progress_cb is None: + return + try: + self._face_detection_progress_cb(payload) + except Exception as e: + logger.warning(f"Failed to emit face detection progress: {e}") + + def _append_smargon_trace(self, *, sample_id: int | None, event: str) -> None: + try: + path = self._smargon_trace_path + path.parent.mkdir(parents=True, exist_ok=True) + + is_new_file = not path.exists() or path.stat().st_size == 0 + pos = self.smargon + sh = pos.sh_mm + + with path.open("a", encoding="utf-8", buffering=1) as f: + if is_new_file: + f.write( + "timestamp,event,sample_id,omega_deg,zoom," + "shx_mm,shy_mm,shz_mm,phi_deg,chi_deg\n" + ) + + f.write( + f"{datetime.now().isoformat(timespec='milliseconds')}," + f"{event}," + f"{'' if sample_id is None else sample_id}," + f"{self.omega:.3f}," + f"{self.zoom:.3f}," + f"{sh.x:.5f}," + f"{sh.y:.5f}," + f"{sh.z:.5f}," + f"{pos.phi_deg:.5f}," + f"{pos.chi_deg:.5f}\n" + ) + f.flush() + + except Exception as e: + logger.warning(f"Failed to append smargon trace: {e}") @property def state(self) -> BeamlineStateEnum: @@ -305,7 +353,22 @@ class AareDAQ: logger.info(f"Mount result: {value}") self.__cfg.current_sample = target #self.__mount_failure_handler(value) - + + def recovery_unmount_sample(self) -> None: + self.__cfg.try_set_busy(timeout=360) + try: + self.__set_state(BeamlineStateEnum.RobotSampleExchange) + self.__devs.tell.check_enable_motion() + self.__devs.tell.wait_not_busy() + self.__devs.tell.set_in_mount_position(True) + self.__devs.tell.unmount(wait=True, timeout=360.0) + self.__cfg.current_sample = None + self.__set_state(BeamlineStateEnum.SampleAlignment) + self.__cfg.state_busy = False + except Exception: + self.__cfg.state_busy = False + raise + @sample.setter def sample(self, target: SampleShortInfo | None): self.__cfg.try_set_busy(timeout=360) @@ -330,7 +393,7 @@ class AareDAQ: if target is not None: if target.db_id is not None: self.__aare.sample_mounted(target) - self.save_screenshot_db(target.db_id, "mounted") + self.save_screenshot_db(target.db_id, f"{target.bd_id}_mounted") @property @@ -843,10 +906,12 @@ class AareDAQ: except Exception as e: logger.error(f"error in face detection sequence {e}") result = { - "samples": None, - "height_fit": None, - "area_fit": None, + "running": False, + "samples": [], + "height_fit": {}, + "area_fit": {}, } + self._emit_face_detection_progress(result) self.__cfg.state_busy = False return result @@ -878,9 +943,6 @@ class AareDAQ: zoom_value = self.__devs.zoom logger.info('face detection sequence') - self.__devs.samcam_settings = SampleCameraSettings( - exposure=zoom_settings[zoom_value].exposure, - gain=zoom_settings[zoom_value].gain) self.__devs.set_zoom(zoom_value, wait=True) boxes: dict[int, tuple[float, float, float, float]] = {} @@ -897,11 +959,18 @@ class AareDAQ: curr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) box_time = time.perf_counter() - m = self.__mlbox.predict(curr_image, filename=None, preferred_class = (3,0)) + m = self.__mlbox.predict(curr_image, filename=None, preferred_class=(3, 0)) logger.info(f"time to predict: {time.perf_counter() - box_time}") if not m or not m.box: logger.info(f"no box found for angle {angle}") + self._emit_face_detection_progress({ + "running": True, + "current_angle_deg": angle, + "samples": fd.get_samples_out(boxes), + "height_fit": {}, + "area_fit": {}, + }) continue cls_id = int(m.cls.value) @@ -915,13 +984,26 @@ class AareDAQ: else: logger.debug(f"ignoring class {cls_id} (pin/crystal) at angle {angle}") - if not boxes: - logger.info("no boxes found") - return {"samples": None, "height_fit": None, "area_fit": None} + self._emit_face_detection_progress({ + "running": True, + "current_angle_deg": angle, + "samples": fd.get_samples_out(boxes), + "height_fit": {}, + "area_fit": {}, + }) + + if not boxes: + logger.info("no boxes found") + result = {"running": False, "samples": [], "height_fit": {}, "area_fit": {}} + self._emit_face_detection_progress(result) + return result best_fit_angle_area, area_params = fd.get_flat_face(boxes, start_angle, end_angle, True) best_fit_angle_height, height_params = fd.get_flat_face(boxes, start_angle, end_angle, False) - fit_results = {"Area":{"angle":best_fit_angle_area, "params":area_params}, "Height": {"angle":best_fit_angle_height, "params":height_params}} + fit_results = { + "Area": {"angle": best_fit_angle_area, "params": area_params}, + "Height": {"angle": best_fit_angle_height, "params": height_params}, + } logger.info(f"best angle by area: {best_fit_angle_area}") logger.info(f"best angle by height: {best_fit_angle_height}") flat_face_angle, best_params, best_name = fd.choose_best_fit(fit_results) @@ -933,24 +1015,29 @@ class AareDAQ: samples_out = fd.get_samples_out(boxes) logger.info(f"face detection sequence done, samples: {samples_out}") - return { + result = { + "running": False, "samples": samples_out, - "height_fit": {"A": height_params["A"], - "B": height_params["B"], - "phi_rad": height_params["phi_rad"], - "C": height_params["C"], - "best_angle_deg": best_fit_angle_height + "height_fit": { + "A": height_params["A"], + "B": height_params["B"], + "phi_rad": height_params["phi_rad"], + "C": height_params["C"], + "best_angle_deg": best_fit_angle_height, }, - "area_fit": {"A": area_params["A"], - "B": area_params["B"], - "phi_rad": area_params["phi_rad"], - "C": height_params["C"], - "best_angle_deg": best_fit_angle_area + "area_fit": { + "A": area_params["A"], + "B": area_params["B"], + "phi_rad": area_params["phi_rad"], + "C": area_params["C"], + "best_angle_deg": best_fit_angle_area, }, } + self._emit_face_detection_progress(result) + return result - def __loop_center_sequence(self, sample_id: int | None = None) -> bool: + def __loop_center_sequence(self, sample_id: int | None = None, trace_all_alc_moves: bool = False) -> bool: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__devs.lamp_light = 2.5 @@ -958,15 +1045,15 @@ class AareDAQ: try: self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter - zoom_settings = self.__cfg.zoom_settings.z + #zoom_settings = self.__cfg.zoom_settings.z - for zoom_iter, zoom_value in enumerate(zoom_settings): + for zoom_iter, zoom_value in enumerate([200]): #exposure = zoom_settings[zoom_value].exposure #gain = zoom_settings[zoom_value].gain max_attempt = 2 attempt = 0 - base_angles = (0, 45, 90) if (zoom_iter % 2 == 0) else (90, 45, 0) + base_angles = (0, 90) if (zoom_iter % 2 == 0) else (90, 0) if sample_id is not None: logger.info(f"submitting to db loop center sequence for sample {sample_id}, zoom={zoom_value}") self.save_screenshot_db(sample_id, f"pre_alc") @@ -1015,6 +1102,11 @@ class AareDAQ: self.__devs.smargon_wait(60) logger.info(f"time to move smargon: {time.perf_counter() - time_to_move_smargon}") if sample_id is not None: + if trace_all_alc_moves: + self._append_smargon_trace( + sample_id=sample_id, + event=f"alc_move_zoom_{zoom_value:.0f}_angle_{angle}" + ) self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{zoom_value:.0f}") if targets_found_this_attempt == 0: @@ -1027,7 +1119,7 @@ class AareDAQ: break if found_flag is not None and found_angle is not None: logger.debug(f"found a target at angle {found_angle} in attempt {attempt + 1}") - base_angles = (found_angle, found_angle + 45, found_angle + 90) + base_angles = (found_angle, found_angle + 45) logger.debug(f"new base angles: {base_angles}") attempt += 1 logger.debug(f"attempt {attempt} of {max_attempt}") @@ -1041,6 +1133,7 @@ class AareDAQ: if sample_id is not None: logger.info(f"sample {sample_id} centered") self.save_screenshot_db(sample_id, f"{sample_id}_centered") + self._append_smargon_trace(sample_id=sample_id, event="alc_success") return True except Exception as e: @@ -1091,6 +1184,18 @@ class AareDAQ: end = time.perf_counter() return end - start + def _default_screenshot_message(self, sample_id: int) -> str: + omega_value = self.omega + zoom_value = self.zoom + samcam = self.samcam_settings + return ( + f"sample_id: {sample_id} " + f"zoom: {zoom_value} " + f"exp:{samcam.exposure} " + f"gain:{samcam.gain} " + f"omega:{omega_value}" + ) + def save_screenshot(self, filename: str): #time.sleep(0.2) # Wait 200 ms to ensure camera image is stable bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) @@ -1101,6 +1206,31 @@ class AareDAQ: bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) self.__aare.upload_image(sample_id, filename, bgr_image) + def send_screenshot_db(self, filename: str | None = None, message: str | None = None) -> None: + sample = self.sample + if sample is None or sample.db_id is None or sample.db_id < 0: + raise ValueError("No sample with a valid sample_id is mounted.") + + sample_id = sample.db_id + bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) + + if filename: + filename = clean_filename(filename) + pgroup = self.__cfg.pgroup + if not pgroup: + raise ValueError("No active pgroup set; cannot save screenshot to photos directory.") + + photos_dir = Path("/sls/mx/data") / pgroup / "raw" / "photos" + photos_dir = photos_dir / str(sample_id) + photos_dir.mkdir(parents=True, exist_ok=True) + photo_path = photos_dir / f"{filename}.jpeg" + cv2.imwrite(str(photo_path), bgr_image) + + upload_name = clean_filename or f"{sample_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + final_message = (message or "").strip() or self._default_screenshot_message(sample_id) + self.__aare.upload_image(sample_id, upload_name, bgr_image, message=final_message) + + @property def sample_spreadsheet(self) -> SampleShortInfoList: return self.__cfg.spreadsheet @@ -1190,7 +1320,6 @@ class AareDAQ: return default_params, "defaults" - def measure(self, sample: SampleShortInfo) -> float: start = time.perf_counter() formatted_date = datetime.now().strftime('%Y%m%d') @@ -1211,7 +1340,7 @@ class AareDAQ: self.__mount(sample) if sample.db_id is not None: self.__aare.sample_mounted(sample) - self.save_screenshot_db(sample.db_id, "mounted") + self.save_screenshot_db(sample.db_id, f"{sample.db_id}_mounted") logger.info(f"mounting done at {time.perf_counter() - start_mount}, total time: {time.perf_counter() - start}") #self.__devs.smargon_pos #self.__devs.aerotech_pos = @@ -1225,7 +1354,9 @@ class AareDAQ: return end - start #raise LoopCenteringFailed logger.info(f"alc done at {time.perf_counter() - start}") - self.__face_detection_sequence() + + result = self.__face_detection_sequence(steps=7, step_size=30) + self._emit_face_detection_progress(result) logger.info(f"face_detection done at {time.perf_counter() - start}") #self.zoom = 500 @@ -1269,7 +1400,6 @@ class AareDAQ: # self.__aare.axc_failed(sample) self.zoom = 1 self.__cfg.zoom_mode = ZoomModeEnum.User - self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom) self.__cfg.state_busy = False except Exception as e: logger.error(f"Error in measure: {e}")