From a00f0b3dff4fc9a413e487aafde4ec7b3d8b41e5 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Thu, 30 Apr 2026 12:44:55 +0200 Subject: [PATCH] DAQ: updated face_detection added operations/face_detection fodler including service and models, new tests of face_detection --- src/aare/daq/daq.py | 420 +++++++++--------- .../daq/operations/face_detection/__init__.py | 11 + .../daq/operations/face_detection/models.py | 24 + .../daq/operations/face_detection/service.py | 213 +++++++++ .../test_face_detection_service.py | 216 +++++++++ tests/unit/daq/test_face_detection.py | 105 +++++ 6 files changed, 783 insertions(+), 206 deletions(-) create mode 100644 src/aare/daq/operations/face_detection/__init__.py create mode 100644 src/aare/daq/operations/face_detection/models.py create mode 100644 src/aare/daq/operations/face_detection/service.py create mode 100644 tests/unit/daq/operations/face_detection/test_face_detection_service.py create mode 100644 tests/unit/daq/test_face_detection.py diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 6a94fb96..2d4a43e3 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -12,7 +12,6 @@ from aareDB import SampleEventType from jfjoch_client.exceptions import NotFoundException from jfjoch_client import ScanResult, ScanResultImagesInner -import aare.common.face_detection as fd from aare.daq import workflows from aare.daq.aaredb import AareWrapper @@ -39,7 +38,7 @@ from aare.common.models import ( PuckLoadedInfo, SampleShortInfoList, AutofocusSettings, DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, ZoomModeEnum, - SimpleScanParameters, MLBoxModel, FluorescenceSpectrumParameterModel, + SimpleScanParameters, FluorescenceSpectrumParameterModel, FluorescenceSpectrumOutputModel, DAQOperation) from aare.common.automation_models import ( AutomationProgress, @@ -50,6 +49,7 @@ from aare.common.automation_models import ( from aare.common.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem, grid_to_image_id from aare.common.rotation_scan import RotationScanRequest, CompletedRotationScan from aare.common.sample_geometry import SampleGeometryModel +from aare.daq.operations.face_detection import FaceDetectionContext, FaceDetectionService from aare.daq.operations.loop_centering import LoopCenteringService, LoopCenteringContext from aare.daq.operations.loop_centering.models import LoopCenteringSettings @@ -282,6 +282,95 @@ class AareDAQ: progress.success = success self._emit_automation_progress(progress) + + #-------------------------------------------- + # Operation Services + #-------------------------------------------- + + def _create_loop_centering_settings(self) -> LoopCenteringSettings: + return LoopCenteringSettings() + + def _create_loop_centering_service(self) -> LoopCenteringService: + settings = self._create_loop_centering_settings() + + return LoopCenteringService( + context=LoopCenteringContext( + cfg=self.__cfg, + devs=self.__devs, + mlbox=self.__mlbox, + settings=settings, + sample_geometry_provider=lambda: self.sample_geometry, + save_screenshot_db=self.save_screenshot_db, + append_smargon_trace=self._append_smargon_trace, + get_predictions=lambda: self.__mlbox.predict_all_best( + overlap_with_pin=settings.overlap_with_pin, + confidence_min=settings.confidence_min, + return_image=True, + return_bundle_meta=True, + ), + ), + logger=logger, + ) + + def _create_face_detection_service(self) -> FaceDetectionService: + return FaceDetectionService( + context=FaceDetectionContext( + cfg=self.__cfg, + devs=self.__devs, + mlbox=self.__mlbox, + sample_geometry_provider=lambda: self.sample_geometry, + emit_progress=self._emit_face_detection_progress, + ), + logger=logger, + ) + # + # def _create_mounting_service(self): + # return MountingService( + # context=MountingContext( + # cfg=self.__cfg, + # devs=self.__devs, + # mount_position=ABR_POS_MOUNT, + # ), + # logger=logger, + # ) + # + # def _create_raster_service(self): + # return RasterService( + # context=RasterContext( + # cfg=self.__cfg, + # devs=self.__devs, + # mlbox=self.__mlbox, + # jfjoch=self.__jfjoch, + # aare=self.__aare, + # sample_provider=lambda: self.sample, + # sample_geometry_provider=lambda: self.sample_geometry, + # status_provider=lambda: self.status, + # set_state=self.__set_state, + # save_screenshot_db=self.save_screenshot_db, + # upload_raster_diffraction_preview=self._upload_raster_diffraction_preview, + # auto_center_line_scan_top_left=self._auto_center_line_scan_top_left, + # ml_bounding_box=self.__ml_bounding_box, + # ), + # logger=logger, + # ) + # + # def _create_rotation_service(self): + # return RotationService( + # context=RotationContext( + # cfg=self.__cfg, + # devs=self.__devs, + # jfjoch=self.__jfjoch, + # aare=self.__aare, + # sample_provider=lambda: self.sample, + # sample_geometry_provider=lambda: self.sample_geometry, + # status_provider=lambda: self.status, + # set_state=self.__set_state, + # save_screenshot_db=self.save_screenshot_db, + # ), + # logger=logger, + # ) + + #-------------------------------------------- # Operation Handlers #-------------------------------------------- @@ -355,6 +444,15 @@ class AareDAQ: """ try: previous_sample = self.sample + if previous_sample is None or previous_sample.db_id is None: + try: + previous_sample = self.sync_current_sample_from_tell( + force=True, + clear_cached_on_empty=False, + ) + except Exception as sync_error: + logger.warning(f"Failed to reconcile previous sample from TELL before mount: {sync_error}") + if previous_sample is not None and previous_sample.db_id is not None: self.__aare.send_sample_event(previous_sample, SampleEventType.UNMOUNTING) self.__set_state(BeamlineStateEnum.RobotSampleExchange) @@ -383,31 +481,6 @@ class AareDAQ: ) return False - def _create_loop_centering_settings(self) -> LoopCenteringSettings: - return LoopCenteringSettings() - - def _create_loop_centering_service(self) -> LoopCenteringService: - settings = self._create_loop_centering_settings() - - return LoopCenteringService( - context=LoopCenteringContext( - cfg=self.__cfg, - devs=self.__devs, - mlbox=self.__mlbox, - settings=settings, - sample_geometry_provider=lambda: self.sample_geometry, - save_screenshot_db=self.save_screenshot_db, - append_smargon_trace=self._append_smargon_trace, - get_predictions=lambda: self.__mlbox.predict_all_best( - overlap_with_pin=settings.overlap_with_pin, - confidence_min=settings.confidence_min, - return_image=True, - return_bundle_meta=True, - ), - ), - logger=logger, - ) - def _execute_loop_centering(self, sample: SampleShortInfo) -> bool: """ Execute loop centering sequence. @@ -458,6 +531,38 @@ class AareDAQ: ) return False + def _execute_face_detection( + self, + *, + steps: int = 14, + step_size: int = 15, + face_min_ratio: float = 0.3, + report_error: bool = True, + ): + """ + Execute face detection sequence through the face detection service. + + Returns: + FaceDetectionResult + """ + self.__set_state(BeamlineStateEnum.SampleAlignment) + + result = self._create_face_detection_service().run( + steps=steps, + step_size=step_size, + face_min_ratio=face_min_ratio, + ) + + if not result.success and report_error: + self._handle_operation_error( + operation=DAQOperation.FACE_CENTERING, + sample=self.sample, + error=result.error or Exception("Face detection failed"), + additional_comment=result.comment, + ) + + return result + def _execute_raster_sequence(self, grid_request: RasterGridRequest, auto_center: bool = False) -> CompletedRasterGrid | None: """ @@ -521,9 +626,19 @@ class AareDAQ: }, ), ) - + else: + if self.sample is not None and self.sample.db_id is not None: + self.__aare.send_sample_event(self.sample, SampleEventType.RASTERINGFAILED) + logger.error( + "Raster sequence returned no result", + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid_request), + ), + ) return result + except JFJochCommunicationError as e: logger.exception( "Raster sequence failed due to JFJoch communication error", @@ -676,7 +791,34 @@ class AareDAQ: location=mounted_address.puck, ) - def sync_current_sample_from_tell(self, force: bool = False) -> SampleShortInfo | None: + def _find_sample_by_mounted_address(self, mounted_address) -> SampleShortInfo | None: + for sample in self.__cfg.spreadsheet.s: + if self._sample_matches_mounted_address(sample, mounted_address): + return sample + + for sample in self.__cfg.reference_tools.s: + if self._sample_matches_mounted_address(sample, mounted_address): + return sample + + return None + + def _placeholder_sample_from_mounted_address(self, mounted_address) -> SampleShortInfo: + return SampleShortInfo( + db_id=-1, + puck_name="", + dewar_name="", + sample_name=f"Mounted sample {mounted_address.puck.segment}{mounted_address.puck.pos}-{mounted_address.pin}", + run_number=0, + user="", + pin=mounted_address.pin, + location=mounted_address.puck, + ) + + def sync_current_sample_from_tell( + self, + force: bool = False, + clear_cached_on_empty: bool = True, + ) -> SampleShortInfo | None: current_sample = self.__cfg.current_sample if current_sample is not None and current_sample.location is None: @@ -691,8 +833,13 @@ class AareDAQ: if mounted_address is None: if current_sample is not None and current_sample.location is not None: - logger.warning("TELL reports no mounted sample; clearing cached current_sample") - self.__cfg.current_sample = None + if clear_cached_on_empty: + logger.warning("TELL reports no mounted sample; clearing cached current_sample") + self.__cfg.current_sample = None + else: + logger.warning( + "TELL reports no mounted sample; keeping cached current_sample to avoid losing context" + ) return self.__cfg.current_sample if self._sample_matches_mounted_address(current_sample, mounted_address): @@ -1166,7 +1313,8 @@ class AareDAQ: }, ), ) - + self.__aare.send_sample_event(self.sample, SampleEventType.RASTERING, + comment=f"Raster at {geom.omega_deg:.1f} deg") r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg") if r is None: @@ -1236,8 +1384,13 @@ class AareDAQ: ) res1 = self.__raster(grid) - grid.omega_deg += 90 + if res1 is None: + self.__aare.send_sample_event(self.sample, SampleEventType.RASTERINGFAILED, + comment=f"No ML Box detected {geom.omega_deg:.1f} deg") + grid.omega_deg += 90 + self.__aare.send_sample_event(self.sample, SampleEventType.RASTERING, + comment=f"Raster at {geom.omega_deg:.1f} deg") self.__devs.aerotech_omega = grid.omega_deg grid.n_x = 1 @@ -1286,6 +1439,11 @@ class AareDAQ: ), ) res2 = self.__raster(grid) + + if res2 is None: + self.__aare.send_sample_event(self.sample, SampleEventType.RASTERINGFAILED, + comment=f"No ML Box detected {geom.omega_deg:.1f} deg") + return CompletedRasterGrid(r=[res1, res2]) else: logger.error( @@ -1505,8 +1663,6 @@ class AareDAQ: try: if self.sample is not None and self.sample.db_id is not None: - self.__aare.send_sample_event(self.sample, SampleEventType.RASTERING, - comment=f"Raster at {request.omega_deg:.1f} deg") self.__aare.create_gridscan_run(self.sample, request, status) self.__jfjoch.wait_till_running(timeout=60.0) @@ -1993,7 +2149,7 @@ class AareDAQ: self.__cfg.state_busy = False raise - def face_detection(self, steps: int = 14, step_size: int = 15, face_min_ratio: float =0.3) -> dict: + def face_detection(self, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3) -> dict: """ Perform a face detection sequence by rotating the sample and using ML to find the flat face. @@ -2001,180 +2157,21 @@ class AareDAQ: steps: Number of rotation steps. Default is 14. step_size: Size of each rotation step in degrees. Default is 15. face_min_ratio: Minimum ratio of loopface count to loop_all count to use loop_face over loop_all. - i.e. if 10 loop_face vs 4 loop_all pick loop_face. if 2 loop_face and 12 loop_all use loop_all. - Default is 0.3. Returns: Dictionary containing face detection results, including found samples and fits. """ self.__cfg.try_set_busy(timeout=360) try: - logger.info("running face detection sequence") - result = self.__face_detection_sequence(steps=steps, step_size=step_size,face_min_ratio=face_min_ratio) - except Exception as e: - logger.error(f"error in face detection sequence {e}") - result = { - "running": False, - "samples": [], - "height_fit": {}, - "area_fit": {}, - } - self._emit_face_detection_progress(result) - self.__cfg.state_busy = False - return result - - def face_detection_centre_correction(self, m:MLBoxModel, tolerance: float = 0.2): - geom = self.sample_geometry - beam_y = geom.beam_location_pxl.y - beam_x = geom.beam_location_pxl.x - - x1 = m.box.top_x - y1 = m.box.top_y - y2 = m.box.bottom_y - - centre_y = y1 + (y2 - y1) / 2 - centre_x = x1 - - if beam_y !=0 and abs(centre_y-beam_y)/abs(beam_y) > tolerance: - coord = geom.picture_to_smargon(Coordinate(x=beam_x, y=centre_y)) - self.__devs.smargon_pos = SmargonCoordinate(sh_mm=coord) - self.__devs.smargon_wait(60) - - return - - #TODO operator function similar to mount, loop_center, raster and rotation - @log_timing(logger, "Face detection sequence") - def __face_detection_sequence(self, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3) -> dict: - self.__set_state(BeamlineStateEnum.SampleAlignment) - self.__devs.lamp_light = 2.5 - self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter - - zoom_value = self.__devs.zoom - logger.info('face detection sequence') - - self.__devs.set_zoom(zoom_value, wait=True) - - boxes_face: dict[int, tuple[float, float, float, float]] = {} - boxes_loop: dict[int, tuple[float, float, float, float]] = {} - curr_angle = int(self.__devs.aerotech_omega) - total_range = steps * step_size + 1 - start_angle = curr_angle if curr_angle + total_range < 720 else 0 - end_angle = curr_angle + total_range - - for angle in range(start_angle, end_angle, step_size): - logger.debug(f'moving to angle: {angle}') - rotate_time = time.perf_counter() - self.__devs.aerotech_omega = angle - log_duration( - logger, - "Completed Aerotech move during face detection", - time.perf_counter() - rotate_time, - extra={"angle_deg": angle}, + result = self._execute_face_detection( + steps=steps, + step_size=step_size, + face_min_ratio=face_min_ratio, + report_error=True, ) - - prediction_result: MLBoxPredictionResult = self.__mlbox.predict( - preferred_class=(3, 0), - return_image=True, - return_bundle_meta=True - ) - m = prediction_result.box - log_ml_bundle_meta( - logger, - f"face_detection_angle_{angle}", - target_point=prediction_result.target_point, - focus=prediction_result.focus, - ) - - 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_face), - "height_fit": {}, - "area_fit": {}, - }) - continue - - cls_id = int(m.cls.value) - x1, y1, x2, y2 = m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y - - self.face_detection_centre_correction(m, tolerance=0.2) - - if cls_id == 3: - boxes_face[angle] = (x1, y1, x2, y2) - logger.info(f"accepted box at angle {angle}, cls={cls_id}, box={(x1, y1, x2, y2)}") - elif cls_id == 0: - boxes_loop[angle] = (x1, y1, x2, y2) - logger.info(f"accepted box at angle {angle}, cls={cls_id}, box={(x1, y1, x2, y2)}") - else: - logger.debug(f"ignoring class {cls_id} at angle {angle}") - - self._emit_face_detection_progress({ - "running": True, - "current_angle_deg": angle, - "samples": fd.get_samples_out(boxes_face), - "height_fit": {}, - "area_fit": {}, - }) - - if not boxes_face and not boxes_loop: - logger.info("no boxes found") - result = {"running": False, "samples": [], "height_fit": {}, "area_fit": {}} - self._emit_face_detection_progress(result) - return result - - total_detections = len(boxes_face) + len(boxes_loop) - face_ratio = len(boxes_face) / total_detections if total_detections > 0 else 0.0 - - if boxes_face and face_ratio >= face_min_ratio: - boxes = boxes_face - logger.info(f"using loop_face boxes ({len(boxes_face)}/{total_detections}, ratio={face_ratio:.2f})") - elif boxes_loop: - boxes = boxes_loop - logger.info( - f"falling back to loop_all boxes ({len(boxes_loop)}/{total_detections}, ratio={1 - face_ratio:.2f})") - else: - boxes = boxes_face - logger.info(f"using loop_face boxes (only source, {len(boxes_face)} entries)") - - 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}, - } - 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) - logger.info(f"best params: {best_params}") - logger.info(f"chosen fit: {best_name}") - logger.info(f"best angle: {flat_face_angle}") - self.__devs.aerotech_omega = flat_face_angle - - samples_out = fd.get_samples_out(boxes) - logger.info(f"face detection sequence done, samples: {samples_out}") - - 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, - }, - "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 + return result.payload + finally: + self.__cfg.state_busy = False def _auto_center_line_scan_top_left( self, @@ -2605,11 +2602,22 @@ class AareDAQ: self._mark_progress_failed(progress, WorkflowStateKind.LOOP_CENTRE, "Centering failed") return self._end_operation(start, DAQOperation.LOOP_CENTERING, error=True) - logger.info(f"Loop Centering done at {time.perf_counter() - start}") - result = self.__face_detection_sequence(steps=7, step_size=30) - self._emit_face_detection_progress(result) + face_detection_result = self._execute_face_detection( + steps=7, + step_size=30, + face_min_ratio=0.3, + report_error=True, + ) + if not face_detection_result.success: + logger.warning( + "Face detection failed during automation; continuing with latest payload", + extra=merge_log_context( + sample_log_context(sample), + {"comment": face_detection_result.comment}, + ), + ) logger.info(f"Face Detection done at {time.perf_counter() - start}") self._mark_progress_success(progress, WorkflowStateKind.LOOP_CENTRE, "Centering complete") diff --git a/src/aare/daq/operations/face_detection/__init__.py b/src/aare/daq/operations/face_detection/__init__.py new file mode 100644 index 00000000..47bc7c2d --- /dev/null +++ b/src/aare/daq/operations/face_detection/__init__.py @@ -0,0 +1,11 @@ +from aare.daq.operations.face_detection.models import ( + FaceDetectionContext, + FaceDetectionResult, +) +from aare.daq.operations.face_detection.service import FaceDetectionService + +__all__ = [ + "FaceDetectionContext", + "FaceDetectionResult", + "FaceDetectionService", +] \ No newline at end of file diff --git a/src/aare/daq/operations/face_detection/models.py b/src/aare/daq/operations/face_detection/models.py new file mode 100644 index 00000000..8aa86ad1 --- /dev/null +++ b/src/aare/daq/operations/face_detection/models.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass +from typing import Callable + +from aare.common.sample_geometry import SampleGeometryModel +from aare.daq.config import BeamlineConfig +from aare.daq.devices import BeamlineDevices +from aare.daq.mlbox import MlBox + + +@dataclass +class FaceDetectionContext: + cfg: BeamlineConfig + devs: BeamlineDevices + mlbox: MlBox + sample_geometry_provider: Callable[[], SampleGeometryModel] + emit_progress: Callable[[dict], None] + + +@dataclass +class FaceDetectionResult: + success: bool + payload: dict + error: Exception | None = None + comment: str | None = None \ No newline at end of file diff --git a/src/aare/daq/operations/face_detection/service.py b/src/aare/daq/operations/face_detection/service.py new file mode 100644 index 00000000..31941523 --- /dev/null +++ b/src/aare/daq/operations/face_detection/service.py @@ -0,0 +1,213 @@ +import time + +import aare.common.face_detection as fd +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.logger_events import log_duration, log_ml_bundle_meta +from aare.common.models import MLBoxModel, ZoomModeEnum +from aare.daq.config import BeamlineStateEnum +from aare.daq.mlbox import MLBoxPredictionResult +from aare.daq.operations.face_detection.models import ( + FaceDetectionContext, + FaceDetectionResult, +) + + +class FaceDetectionService: + def __init__(self, *, context: FaceDetectionContext, logger): + self.ctx = context + self.logger = logger + + def _emit_running_progress( + self, + *, + angle: int, + boxes_face: dict[int, tuple[float, float, float, float]], + ) -> None: + self.ctx.emit_progress( + { + "running": True, + "current_angle_deg": angle, + "samples": fd.get_samples_out(boxes_face), + "height_fit": {}, + "area_fit": {}, + } + ) + + def _emit_empty_result(self) -> dict: + payload = { + "running": False, + "samples": [], + "height_fit": {}, + "area_fit": {}, + } + self.ctx.emit_progress(payload) + return payload + + def _centre_correction(self, model: MLBoxModel, tolerance: float = 0.2) -> None: + geom = self.ctx.sample_geometry_provider() + beam_y = geom.beam_location_pxl.y + beam_x = geom.beam_location_pxl.x + + x1 = model.box.top_x + y1 = model.box.top_y + y2 = model.box.bottom_y + + centre_y = y1 + (y2 - y1) / 2 + + if beam_y != 0 and abs(centre_y - beam_y) / abs(beam_y) > tolerance: + coord = geom.picture_to_smargon(Coordinate(x=beam_x, y=centre_y)) + self.ctx.devs.smargon_pos = SmargonCoordinate(sh_mm=coord) + self.ctx.devs.smargon_wait(60) + + def run( + self, + *, + steps: int = 14, + step_size: int = 15, + face_min_ratio: float = 0.3, + ) -> FaceDetectionResult: + try: + self.ctx.cfg.zoom_mode = ZoomModeEnum.LoopCenter + self.ctx.devs.lamp_light = 2.5 + + zoom_value = self.ctx.devs.zoom + self.logger.info("face detection sequence") + self.ctx.devs.set_zoom(zoom_value, wait=True) + + boxes_face: dict[int, tuple[float, float, float, float]] = {} + boxes_loop: dict[int, tuple[float, float, float, float]] = {} + + curr_angle = int(self.ctx.devs.aerotech_omega) + total_range = steps * step_size + 1 + start_angle = curr_angle if curr_angle + total_range < 720 else 0 + end_angle = curr_angle + total_range + + for angle in range(start_angle, end_angle, step_size): + self.logger.debug(f"moving to angle: {angle}") + rotate_time = time.perf_counter() + self.ctx.devs.aerotech_omega = angle + log_duration( + self.logger, + "Completed Aerotech move during face detection", + time.perf_counter() - rotate_time, + extra={"angle_deg": angle}, + ) + + prediction_result: MLBoxPredictionResult = self.ctx.mlbox.predict( + preferred_class=(3, 0), + return_image=True, + return_bundle_meta=True, + ) + model = prediction_result.box + log_ml_bundle_meta( + self.logger, + f"face_detection_angle_{angle}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) + + if not model or not model.box: + self.logger.info(f"no box found for angle {angle}") + self._emit_running_progress(angle=angle, boxes_face=boxes_face) + continue + + cls_id = int(model.cls.value) + x1, y1, x2, y2 = ( + model.box.top_x, + model.box.top_y, + model.box.bottom_x, + model.box.bottom_y, + ) + + self._centre_correction(model, tolerance=0.2) + + if cls_id == 3: + boxes_face[angle] = (x1, y1, x2, y2) + self.logger.info( + f"accepted box at angle {angle}, cls={cls_id}, box={(x1, y1, x2, y2)}" + ) + elif cls_id == 0: + boxes_loop[angle] = (x1, y1, x2, y2) + self.logger.info( + f"accepted box at angle {angle}, cls={cls_id}, box={(x1, y1, x2, y2)}" + ) + else: + self.logger.debug(f"ignoring class {cls_id} at angle {angle}") + + self._emit_running_progress(angle=angle, boxes_face=boxes_face) + + if not boxes_face and not boxes_loop: + self.logger.info("no boxes found") + payload = self._emit_empty_result() + return FaceDetectionResult(success=True, payload=payload) + + total_detections = len(boxes_face) + len(boxes_loop) + face_ratio = len(boxes_face) / total_detections if total_detections > 0 else 0.0 + + if boxes_face and face_ratio >= face_min_ratio: + boxes = boxes_face + self.logger.info( + f"using loop_face boxes ({len(boxes_face)}/{total_detections}, ratio={face_ratio:.2f})" + ) + elif boxes_loop: + boxes = boxes_loop + self.logger.info( + f"falling back to loop_all boxes ({len(boxes_loop)}/{total_detections}, ratio={1 - face_ratio:.2f})" + ) + else: + boxes = boxes_face + self.logger.info(f"using loop_face boxes (only source, {len(boxes_face)} entries)") + + 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}, + } + + self.logger.info(f"best angle by area: {best_fit_angle_area}") + self.logger.info(f"best angle by height: {best_fit_angle_height}") + + flat_face_angle, best_params, best_name = fd.choose_best_fit(fit_results) + + self.logger.info(f"best params: {best_params}") + self.logger.info(f"chosen fit: {best_name}") + self.logger.info(f"best angle: {flat_face_angle}") + + self.ctx.devs.aerotech_omega = flat_face_angle + + samples_out = fd.get_samples_out(boxes) + self.logger.info(f"face detection sequence done, samples: {samples_out}") + + payload = { + "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, + }, + "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.ctx.emit_progress(payload) + return FaceDetectionResult(success=True, payload=payload) + + except Exception as e: + self.logger.error(f"error in face detection sequence {e}") + payload = self._emit_empty_result() + return FaceDetectionResult( + success=False, + payload=payload, + error=e, + comment="Face detection sequence failed", + ) + finally: + self.ctx.cfg.zoom_mode = ZoomModeEnum.User \ No newline at end of file diff --git a/tests/unit/daq/operations/face_detection/test_face_detection_service.py b/tests/unit/daq/operations/face_detection/test_face_detection_service.py new file mode 100644 index 00000000..15d43147 --- /dev/null +++ b/tests/unit/daq/operations/face_detection/test_face_detection_service.py @@ -0,0 +1,216 @@ +import types + +import pytest + +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.models import BoundingBoxModel, MLBoxModel, MLBoxType, ZoomModeEnum +from aare.daq.operations.face_detection.models import ( + FaceDetectionContext, + FaceDetectionResult, +) +from aare.daq.operations.face_detection.service import FaceDetectionService + + +class DummyLogger: + def info(self, *args, **kwargs): + pass + + def debug(self, *args, **kwargs): + pass + + def error(self, *args, **kwargs): + pass + + def log(self, *args, **kwargs): + pass + + +class DummyGeometry: + def __init__(self): + self.beam_location_pxl = Coordinate(x=100.0, y=100.0, z=0.0) + + def picture_to_smargon(self, coord: Coordinate) -> Coordinate: + return Coordinate(x=coord.x / 100.0, y=coord.y / 100.0, z=0.0) + + +def _box(cls: MLBoxType, x1: float, y1: float, x2: float, y2: float) -> MLBoxModel: + return MLBoxModel( + cls=cls, + box=BoundingBoxModel( + top_x=x1, + top_y=y1, + bottom_x=x2, + bottom_y=y2, + ), + conf=0.95, + ) + + +@pytest.fixture +def context(): + progress_events = [] + + devs = types.SimpleNamespace( + lamp_light=0.0, + zoom=200.0, + aerotech_omega=0.0, + smargon_pos=None, + set_zoom=lambda zoom_value, wait=True: None, + smargon_wait=lambda timeout=60: None, + ) + + cfg = types.SimpleNamespace(zoom_mode=ZoomModeEnum.User) + + ctx = FaceDetectionContext( + cfg=cfg, + devs=devs, + mlbox=types.SimpleNamespace( + predict=lambda **kwargs: types.SimpleNamespace( + box=None, + target_point=None, + focus=None, + ) + ), + sample_geometry_provider=lambda: DummyGeometry(), + emit_progress=lambda payload: progress_events.append(payload), + ) + ctx._progress_events = progress_events + return ctx + + +def test_service_returns_empty_payload_when_no_boxes(monkeypatch, context): + service = FaceDetectionService(context=context, logger=DummyLogger()) + + monkeypatch.setattr( + context.mlbox, + "predict", + lambda **kwargs: types.SimpleNamespace( + box=None, + target_point=None, + focus=None, + ), + ) + + result = service.run(steps=2, step_size=30) + + assert isinstance(result, FaceDetectionResult) + assert result.success is True + assert result.payload["running"] is False + assert result.payload["samples"] == [] + assert result.payload["height_fit"] == {} + assert result.payload["area_fit"] == {} + assert context.cfg.zoom_mode == ZoomModeEnum.User + assert len(context._progress_events) >= 1 + assert context._progress_events[-1]["running"] is False + + +def test_service_prefers_face_boxes_when_ratio_is_high(monkeypatch, context): + service = FaceDetectionService(context=context, logger=DummyLogger()) + + predictions = iter( + [ + types.SimpleNamespace(box=_box(MLBoxType.LOOP_FACE, 10, 20, 30, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_FACE, 12, 20, 32, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_ALL, 14, 20, 34, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_FACE, 16, 20, 36, 40), target_point=None, focus=None), + ] + ) + + monkeypatch.setattr(context.mlbox, "predict", lambda **kwargs: next(predictions)) + + monkeypatch.setattr( + "aare.daq.operations.face_detection.service.fd.get_flat_face", + lambda boxes, start_angle, end_angle, use_area: ( + 45 if use_area else 50, + {"A": 1.0, "B": 2.0, "phi_rad": 0.1, "C": 3.0}, + ), + ) + monkeypatch.setattr( + "aare.daq.operations.face_detection.service.fd.choose_best_fit", + lambda fit_results: (47, {"A": 1.0}, "Area"), + ) + monkeypatch.setattr( + "aare.daq.operations.face_detection.service.fd.get_samples_out", + lambda boxes: [{"angle": angle, "box": box} for angle, box in sorted(boxes.items())], + ) + + result = service.run(steps=3, step_size=30, face_min_ratio=0.5) + + assert result.success is True + assert result.payload["running"] is False + assert result.payload["height_fit"]["best_angle_deg"] == 50 + assert result.payload["area_fit"]["best_angle_deg"] == 45 + assert context.devs.aerotech_omega == 47 + assert [sample["angle"] for sample in result.payload["samples"]] == [0, 30, 90] + assert context._progress_events[-1]["running"] is False + + +def test_service_falls_back_to_loop_all_when_face_ratio_is_low(monkeypatch, context): + service = FaceDetectionService(context=context, logger=DummyLogger()) + + predictions = iter( + [ + types.SimpleNamespace(box=_box(MLBoxType.LOOP_FACE, 10, 20, 30, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_ALL, 11, 20, 31, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_ALL, 12, 20, 32, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_ALL, 13, 20, 33, 40), target_point=None, focus=None), + types.SimpleNamespace(box=_box(MLBoxType.LOOP_ALL, 14, 20, 34, 40), target_point=None, focus=None), + ] + ) + + monkeypatch.setattr(context.mlbox, "predict", lambda **kwargs: next(predictions)) + + captured = {} + + def fake_get_flat_face(boxes, start_angle, end_angle, use_area): + captured["boxes_used"] = dict(boxes) + return 60, {"A": 1.0, "B": 2.0, "phi_rad": 0.2, "C": 4.0} + + monkeypatch.setattr("aare.daq.operations.face_detection.service.fd.get_flat_face", fake_get_flat_face) + monkeypatch.setattr( + "aare.daq.operations.face_detection.service.fd.choose_best_fit", + lambda fit_results: (60, {"A": 1.0}, "Height"), + ) + monkeypatch.setattr( + "aare.daq.operations.face_detection.service.fd.get_samples_out", + lambda boxes: [{"angle": angle} for angle in sorted(boxes.keys())], + ) + + result = service.run(steps=4, step_size=30, face_min_ratio=0.5) + + assert result.success is True + assert len(captured["boxes_used"]) == 4 + assert sorted(captured["boxes_used"].keys()) == [30, 60, 90, 120] + assert context.devs.aerotech_omega == 60 + assert context._progress_events[-1]["running"] is False + + +def test_service_applies_centre_correction_when_target_is_far_from_beam(monkeypatch, context): + service = FaceDetectionService(context=context, logger=DummyLogger()) + + model = _box(MLBoxType.LOOP_FACE, 40, 160, 80, 200) + + service._centre_correction(model, tolerance=0.2) + + assert isinstance(context.devs.smargon_pos, SmargonCoordinate) + assert context.devs.smargon_pos.sh_mm.x == pytest.approx(1.0) + assert context.devs.smargon_pos.sh_mm.y == pytest.approx(1.8) + + +def test_service_returns_failed_result_when_prediction_raises(monkeypatch, context): + service = FaceDetectionService(context=context, logger=DummyLogger()) + + monkeypatch.setattr( + context.mlbox, + "predict", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("prediction failed")), + ) + + result = service.run(steps=1, step_size=30) + + assert result.success is False + assert isinstance(result.error, RuntimeError) + assert result.comment == "Face detection sequence failed" + assert result.payload["running"] is False + assert context.cfg.zoom_mode == ZoomModeEnum.User + assert context._progress_events[-1]["running"] is False \ No newline at end of file diff --git a/tests/unit/daq/test_face_detection.py b/tests/unit/daq/test_face_detection.py new file mode 100644 index 00000000..2755f323 --- /dev/null +++ b/tests/unit/daq/test_face_detection.py @@ -0,0 +1,105 @@ +import types + +from aare.common.models import DAQOperation +from aare.daq.operations.face_detection.models import FaceDetectionResult + + +def test_execute_face_detection_reports_failure(monkeypatch): + from aare.daq.daq import AareDAQ + + daq = object.__new__(AareDAQ) + + calls = {"set_state": [], "handle_error": []} + + monkeypatch.setattr( + daq, + "_AareDAQ__set_state", + lambda state: calls["set_state"].append(state), + ) + monkeypatch.setattr( + daq, + "_create_face_detection_service", + lambda: types.SimpleNamespace( + run=lambda **kwargs: FaceDetectionResult( + success=False, + payload={"running": False, "samples": [], "height_fit": {}, "area_fit": {}}, + error=RuntimeError("fd failed"), + comment="face detection sequence failed", + ) + ), + ) + monkeypatch.setattr( + daq, + "_handle_operation_error", + lambda **kwargs: calls["handle_error"].append(kwargs), + ) + monkeypatch.setattr( + type(daq), + "sample", + property(lambda self: types.SimpleNamespace(sample_name="sample-1", db_id=1)), + ) + + result = daq._execute_face_detection(steps=7, step_size=30, report_error=True) + + assert result.success is False + assert len(calls["set_state"]) == 1 + assert len(calls["handle_error"]) == 1 + assert calls["handle_error"][0]["operation"] == DAQOperation.FACE_CENTERING + + +def test_execute_face_detection_can_skip_error_reporting(monkeypatch): + from aare.daq.daq import AareDAQ + + daq = object.__new__(AareDAQ) + + handle_error_calls = [] + + monkeypatch.setattr(daq, "_AareDAQ__set_state", lambda state: None) + monkeypatch.setattr( + daq, + "_create_face_detection_service", + lambda: types.SimpleNamespace( + run=lambda **kwargs: FaceDetectionResult( + success=False, + payload={"running": False, "samples": [], "height_fit": {}, "area_fit": {}}, + error=RuntimeError("fd failed"), + comment="face detection sequence failed", + ) + ), + ) + monkeypatch.setattr( + daq, + "_handle_operation_error", + lambda **kwargs: handle_error_calls.append(kwargs), + ) + + result = daq._execute_face_detection(report_error=False) + + assert result.success is False + assert handle_error_calls == [] + + +def test_public_face_detection_uses_execute_face_detection(monkeypatch): + from aare.daq.daq import AareDAQ + + daq = object.__new__(AareDAQ) + cfg = types.SimpleNamespace( + try_set_busy=lambda timeout=360: None, + state_busy=False, + ) + setattr(daq, "_AareDAQ__cfg", cfg) + + monkeypatch.setattr( + daq, + "_execute_face_detection", + lambda **kwargs: FaceDetectionResult( + success=True, + payload={"running": False, "samples": [{"angle": 45}], "height_fit": {}, "area_fit": {}}, + ), + ) + + result = daq.face_detection(steps=7, step_size=30) + + assert result["running"] is False + assert result["samples"] == [{"angle": 45}] + assert cfg.state_busy is False \ No newline at end of file