diff --git a/src/aare/common/config/x06da.yaml b/src/aare/common/config/x06da.yaml index 2bf17cd8..95d8bbaa 100644 --- a/src/aare/common/config/x06da.yaml +++ b/src/aare/common/config/x06da.yaml @@ -30,6 +30,8 @@ daq: lens_magnification: 10 #or 5 currently default_detector_distance_minimum: 86 default_detector_distance_maximum: 900 + zoom_min: 1 # camera zoom travel limits (used to clamp auto-center zoom-to-fit) + zoom_max: 1000 db: aaredb_url: "https://mx-aaredb-dmz-01.psi.ch/dispatcher" @@ -50,3 +52,7 @@ daq: detector_limit_modifier: 2.0 maximum_flux: 4e11 + + auto_raster: + loop_face_padding_fraction: 0.15 # pad loop_face raster boxes by 15% per side + include_crystal: false # extend the grid to cover crystals outside the loop box diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index c1327e42..c6c8d687 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -625,6 +625,28 @@ class BeamlineConfig: lens_magnification = DEFAULT_LENS_MAGNIFICATION return base_pixel_in_mm * (DEFAULT_LENS_MAGNIFICATION / lens_magnification) + def zoom_for_pixel_to_mm(self, target_pixel_in_mm: float) -> float: + """Inverse of :meth:`pixel_to_mm`: the zoom at which one pixel covers + ``target_pixel_in_mm`` millimetres. + + pixel_to_mm(z) = lens_factor / (b * exp(a*z)) => + z = ln(lens_factor / (b * target)) / a + """ + if target_pixel_in_mm <= 0: + raise ValueError(f"target_pixel_in_mm must be > 0, got {target_pixel_in_mm}") + cfg = self.settings + a = cfg.camera_translation_factor_a + b = cfg.camera_translation_factor_b + lens_magnification = cfg_get("daq.hardware.lens_magnification", DEFAULT_LENS_MAGNIFICATION) + try: + lens_magnification = float(lens_magnification) + except (TypeError, ValueError): + lens_magnification = DEFAULT_LENS_MAGNIFICATION + if lens_magnification <= 0: + lens_magnification = DEFAULT_LENS_MAGNIFICATION + lens_factor = DEFAULT_LENS_MAGNIFICATION / lens_magnification + return float(np.log(lens_factor / (b * target_pixel_in_mm)) / a) + @property def beam_center(self) -> Tuple[float, float]: tmp_x = self.__client.get(f"{self.__bl}:beam_center_x") diff --git a/src/aare/daq/operations/common/ml_bounding_box.py b/src/aare/daq/operations/common/ml_bounding_box.py index 19d4df69..0428f225 100644 --- a/src/aare/daq/operations/common/ml_bounding_box.py +++ b/src/aare/daq/operations/common/ml_bounding_box.py @@ -1,9 +1,11 @@ import time +from dataclasses import dataclass from math import ceil, floor from typing import Callable import cv2 +from aare.common.beamline import cfg_get from aare.common.coordinate import Coordinate, SmargonCoordinate from aare.common.exception_handler import AutoRasterSampleSkipped from aare.common.logger_events import ( @@ -12,10 +14,39 @@ from aare.common.logger_events import ( merge_log_context, sample_log_context, ) -from aare.common.models import SampleShortInfo +from aare.common.models import SampleShortInfo, MLBoxType from aare.common.raster_grid import RasterGridRequest from aare.common.sample_geometry import SampleGeometryModel -from aare.daq.mlbox import MLBoxPredictionResult, MlBox +from aare.daq.mlbox import MLBoxPredictionResult, MLBoxPredictionsResult, MlBox + + +BoxTuple = tuple[float, float, float, float] + + +@dataclass +class MLRasterPlan: + """Result of an auto-center ML detection: the raster grid request plus the + raw loop boxes (at the current zoom) needed to drive zoom-to-fit.""" + grid_request: RasterGridRequest + loop_all_box: BoxTuple | None + loop_face_box: BoxTuple | None + image_width: int | None + image_height: int | None + + +def _box_tuple(model) -> BoxTuple | None: + if model is None or model.box is None: + return None + return (model.box.top_x, model.box.top_y, model.box.bottom_x, model.box.bottom_y) + + +def _box_extends_beyond(inner: BoxTuple, outer: BoxTuple) -> bool: + return (inner[0] < outer[0] or inner[1] < outer[1] + or inner[2] > outer[2] or inner[3] > outer[3]) + + +def _box_union(a: BoxTuple, b: BoxTuple) -> BoxTuple: + return (min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3])) def scale_auto_raster_grid( @@ -133,6 +164,32 @@ def get_ml_bounding_box( ), ) + return _box_to_raster_request( + x1=x1, y1=y1, x2=x2, y2=y2, + sample=sample, + sample_geometry=geom, + logger=logger, + filename=filename, + sample_id=sample_id, + max_images=max_images, + min_cell_size_mm=min_cell_size_mm, + skip_if_exceed_max_image_threshold=skip_if_exceed_max_image_threshold, + ) + + +def _box_to_raster_request( + *, + x1: float, y1: float, x2: float, y2: float, + sample: SampleShortInfo | None, + sample_geometry: SampleGeometryModel, + logger, + filename: str | None, + sample_id: int | None, + max_images: int, + min_cell_size_mm: float, + skip_if_exceed_max_image_threshold: bool, +) -> RasterGridRequest: + geom = sample_geometry start_coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1)) grid_size = Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.8) n_x = max(1, abs(ceil((x2 - x1) * geom.pixel_in_mm / grid_size.x))) @@ -211,4 +268,109 @@ def get_ml_bounding_box( n_y=n_y, grid_size_mm=grid_size, omega_deg=geom.omega_deg, + ) + + +def build_ml_raster_plan( + *, + mlbox: MlBox, + sample: SampleShortInfo | None, + sample_geometry: SampleGeometryModel, + upload_image: Callable[[int | None, str, object], None], + logger, + filename: str | None = None, + max_images: int, + min_cell_size_mm: float, + skip_if_exceed_max_image_threshold: bool, +) -> MLRasterPlan | None: + """Like :func:`get_ml_bounding_box`, but used by the auto-center path. + + It runs a single prediction and returns the raster grid request together + with the raw ``loop_all`` / ``loop_face`` boxes (so the caller can drive + zoom-to-fit). The grid box is loop_face (else loop_all); a loop_face box is + padded by ``daq.auto_raster.loop_face_padding_fraction`` and, when + ``daq.auto_raster.include_crystal`` is enabled, the box is extended to the + union of itself and any detected crystal boxes that lie beyond it. + """ + time.sleep(0.2) + sample_id = getattr(sample, "db_id", None) + + result: MLBoxPredictionsResult = mlbox.predict_all_best( + return_image=True, + return_bundle_meta=True, + ) + predictions = result.predictions + bundle_image = result.image + + log_ml_bundle_meta( + logger, + f"ml_raster_plan:{filename or 'unnamed'}", + target_point=result.target_point, + focus=result.focus, + ) + + loop_all = predictions.get_best_for_class(MLBoxType.LOOP_ALL) if predictions else None + loop_face = predictions.get_best_for_class(MLBoxType.LOOP_FACE) if predictions else None + + # Grid box: prefer loop_face, else loop_all (matches the legacy (3, 0) order). + grid_model = loop_face if loop_face is not None else loop_all + if grid_model is None: + logger.warning( + "ML raster plan returned no loop detection", + extra={"sample_id": sample_id, "ml_image_name": filename, + "target_point": result.target_point}, + ) + if filename is not None and bundle_image is not None: + upload_image(sample_id, f"{filename}_no_detection", bundle_image) + return None + + x1, y1, x2, y2 = (grid_model.box.top_x, grid_model.box.top_y, + grid_model.box.bottom_x, grid_model.box.bottom_y) + is_loop_face = grid_model is loop_face + + # Pad loop_face boxes a little so the grid isn't cut tight to the face. + if is_loop_face: + frac = float(cfg_get("daq.auto_raster.loop_face_padding_fraction", 0.15)) + pad_x = (x2 - x1) * frac + pad_y = (y2 - y1) * frac + x1, y1, x2, y2 = x1 - pad_x, y1 - pad_y, x2 + pad_x, y2 + pad_y + + # Optionally extend the grid to cover crystals detected outside the loop box. + if cfg_get("daq.auto_raster.include_crystal", False) and predictions is not None: + for crystal in predictions.get_models_for_class(MLBoxType.CRYSTAL): + cbox = (crystal.box.top_x, crystal.box.top_y, crystal.box.bottom_x, crystal.box.bottom_y) + if _box_extends_beyond(cbox, (x1, y1, x2, y2)): + x1, y1, x2, y2 = _box_union((x1, y1, x2, y2), cbox) + logger.info( + "Extended ML raster grid to include crystal outside the loop box", + extra={"sample_id": sample_id, "ml_image_name": filename, + "crystal_box": cbox}, + ) + + if filename is not None and bundle_image is not None: + annotated_image = bundle_image.copy() + cv2.rectangle(annotated_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) + upload_image(sample_id, filename, annotated_image) + + grid_request = _box_to_raster_request( + x1=x1, y1=y1, x2=x2, y2=y2, + sample=sample, + sample_geometry=sample_geometry, + logger=logger, + filename=filename, + sample_id=sample_id, + max_images=max_images, + min_cell_size_mm=min_cell_size_mm, + skip_if_exceed_max_image_threshold=skip_if_exceed_max_image_threshold, + ) + + image_height = int(bundle_image.shape[0]) if bundle_image is not None else None + image_width = int(bundle_image.shape[1]) if bundle_image is not None else None + + return MLRasterPlan( + grid_request=grid_request, + loop_all_box=_box_tuple(loop_all), + loop_face_box=_box_tuple(loop_face), + image_width=image_width, + image_height=image_height, ) \ No newline at end of file diff --git a/src/aare/daq/operations/raster/service.py b/src/aare/daq/operations/raster/service.py index 8b5804b3..2496abd0 100644 --- a/src/aare/daq/operations/raster/service.py +++ b/src/aare/daq/operations/raster/service.py @@ -16,12 +16,18 @@ from aare.common.logger_events import ( raster_request_log_context, sample_log_context, ) +from aare.common.beamline import cfg_get from aare.common.models import BeamlineStateEnum from aare.common.raster_grid import CompletedRasterGrid, CompletedRasterGridElem, RasterGridRequest, grid_to_image_id from aare.common.simulate_raster import generate_no_beam_scan_result from aare.daq.mlbox import MLBoxPredictionResult -from aare.daq.operations.common.ml_bounding_box import get_ml_bounding_box +from aare.daq.operations.common.ml_bounding_box import ( + MLRasterPlan, + build_ml_raster_plan, + get_ml_bounding_box, +) from aare.daq.operations.raster.models import RasterBoundingBoxResult, RasterContext +from aare.devices.area_detector import AutoEnum class RasterService: @@ -120,6 +126,85 @@ class RasterService: skip_if_exceed_max_image_threshold=self.ctx.settings.auto_raster_skip_if_exceed_max_image_threshold, ) + def ml_raster_plan( + self, + sample_id: int | None = None, + filename: str | None = None, + ) -> MLRasterPlan | None: + return build_ml_raster_plan( + mlbox=self.ctx.deps.mlbox, + sample=self.ctx.sample, + sample_geometry=self.ctx.sample_geometry, + upload_image=self.ctx.deps.aare.upload_image, + logger=self.logger, + filename=filename, + max_images=self.ctx.settings.auto_raster_max_images, + min_cell_size_mm=self.ctx.settings.auto_raster_min_cell_size_mm, + skip_if_exceed_max_image_threshold=self.ctx.settings.auto_raster_skip_if_exceed_max_image_threshold, + ) + + @staticmethod + def _box_touches_frame_edge(box, width: int | None, height: int | None, margin: int = 2) -> bool: + if width is None or height is None: + return False + x1, y1, x2, y2 = box + return x1 <= margin or y1 <= margin or x2 >= width - margin or y2 >= height - margin + + def _zoom_to_fit_box(self, plan: MLRasterPlan) -> tuple | None: + """Pick the box that drives zoom-to-fit: loop_all if detected and not + clipped at the frame edge, otherwise loop_face.""" + if plan.loop_all_box is not None and not self._box_touches_frame_edge( + plan.loop_all_box, plan.image_width, plan.image_height + ): + return plan.loop_all_box + return plan.loop_face_box + + def _apply_zoom_to_fit(self, plan: MLRasterPlan) -> None: + """Zoom so the chosen loop box covers at most 1/3 of the frame.""" + box = self._zoom_to_fit_box(plan) + if box is None or not plan.image_width or not plan.image_height: + return + + x1, y1, x2, y2 = box + box_px = max(x2 - x1, y2 - y1) + if box_px <= 0: + return + screen_px = min(plan.image_width, plan.image_height) + + current_zoom = self.ctx.deps.devs.zoom + pixel_in_mm_now = self.ctx.deps.cfg.pixel_to_mm(current_zoom) + box_mm = box_px * pixel_in_mm_now + # Target mm-per-pixel so the box spans exactly screen/3. + target_pixel_in_mm = 3.0 * box_mm / screen_px + try: + target_zoom = self.ctx.deps.cfg.zoom_for_pixel_to_mm(target_pixel_in_mm) + except (ValueError, ZeroDivisionError) as e: + self.logger.warning(f"Skipping auto-center zoom-to-fit: {e}") + return + + zoom_min = float(cfg_get("daq.hardware.zoom_min", 1.0)) + zoom_max = float(cfg_get("daq.hardware.zoom_max", 1000.0)) + target_zoom = max(zoom_min, min(zoom_max, target_zoom)) + + self.logger.info( + "Auto-center zoom-to-fit", + extra=merge_log_context( + sample_log_context(self.ctx.sample), + { + "box_px": box_px, + "screen_px": screen_px, + "current_zoom": current_zoom, + "target_zoom": target_zoom, + "used_loop_all": box is plan.loop_all_box, + }, + ), + ) + + self.ctx.deps.devs.samcam_auto(AutoEnum.AUTO) + self.ctx.deps.devs.set_zoom(target_zoom, wait=True) + time.sleep(0.2) + self.ctx.deps.devs.samcam_auto(AutoEnum.ONCE) + def auto_center_line_scan_top_left( self, *, @@ -598,9 +683,9 @@ class RasterService: comment=f"Raster at {geom.omega_deg:.1f} deg", ) - r = self.ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg") + plan = self.ml_raster_plan(sample.db_id, f"ml_{geom.omega_deg:.2f}deg") - if r is None: + if plan is None: self.logger.warning( "No ML bounding box found at primary angle during auto-center raster", extra=merge_log_context( @@ -613,9 +698,9 @@ class RasterService: ) self.ctx.deps.devs.aerotech_omega = geom.omega_deg + 90.0 time.sleep(0.2) - r = self.ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg + 90.0:.2f}deg") + plan = self.ml_raster_plan(sample.db_id, f"ml_{geom.omega_deg + 90.0:.2f}deg") - if r is None: + if plan is None: self.logger.error( "Auto-center raster aborted because no ML bounding box was found at either angle", extra=merge_log_context( @@ -629,6 +714,11 @@ class RasterService: ) return None + # Zoom so the loop fills at most 1/3 of the frame before the first raster. + # The grid was built from this same detection (no re-predict). + self._apply_zoom_to_fit(plan) + r = plan.grid_request + self.logger.info( "ML bounding box found for auto-center raster", extra=merge_log_context( diff --git a/tests/unit/daq/operations/test_ml_raster_plan.py b/tests/unit/daq/operations/test_ml_raster_plan.py new file mode 100644 index 00000000..72a53ec1 --- /dev/null +++ b/tests/unit/daq/operations/test_ml_raster_plan.py @@ -0,0 +1,131 @@ +import logging +import types + +import numpy as np +import pytest + +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.models import MLBoxType, MLOutputModel +from aare.common.sample_geometry import SampleGeometryModel +from aare.daq.mlbox import MLBoxPredictionsResult +from aare.daq.operations.common import ml_bounding_box as mlb +from aare.daq.operations.common.ml_bounding_box import ( + _box_extends_beyond, + _box_union, + build_ml_raster_plan, +) + +logger = logging.getLogger("test_ml_raster_plan") + + +def _geom() -> SampleGeometryModel: + return SampleGeometryModel( + beam_location_pxl=Coordinate(x=500, y=500), + pixel_in_mm=0.001, + aerotech=Coordinate(x=0, y=0), + aerotech_meas=Coordinate(x=0, y=0), + smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0, z=0), phi_deg=0, chi_deg=0), + omega_deg=0, + beam_size_mm=Coordinate(x=0.01, y=0.01), + ) + + +def _fake_mlbox(loop_all=None, loop_face=None, crystals=()): + preds = MLOutputModel() + if loop_all is not None: + preds.add_box(MLBoxType.LOOP_ALL, loop_all, 0.9) + if loop_face is not None: + preds.add_box(MLBoxType.LOOP_FACE, loop_face, 0.8) + for c in crystals: + preds.add_box(MLBoxType.CRYSTAL, c, 0.7) + result = MLBoxPredictionsResult( + predictions=preds if preds.boxes else None, + image=np.zeros((1000, 1000, 3), dtype=np.uint8), + target_point=None, + focus=None, + ) + return types.SimpleNamespace(predict_all_best=lambda **kwargs: result) + + +def _plan(mlbox): + return build_ml_raster_plan( + mlbox=mlbox, + sample=None, + sample_geometry=_geom(), + upload_image=lambda *a, **k: None, + logger=logger, + filename=None, + max_images=100000, + min_cell_size_mm=0.0001, + skip_if_exceed_max_image_threshold=False, + ) + + +def test_box_helpers(): + assert _box_union((1, 1, 3, 3), (2, 0, 5, 4)) == (1, 0, 5, 4) + assert _box_extends_beyond((350, 150, 500, 300), (100, 100, 400, 400)) is True + assert _box_extends_beyond((150, 150, 300, 300), (100, 100, 400, 400)) is False + + +def test_plan_returns_loop_boxes_and_prefers_loop_face(): + plan = _plan(_fake_mlbox(loop_all=(100, 100, 400, 400), loop_face=(150, 150, 300, 300))) + assert plan is not None + assert plan.loop_all_box == (100, 100, 400, 400) + assert plan.loop_face_box == (150, 150, 300, 300) + assert plan.image_width == 1000 and plan.image_height == 1000 + assert plan.grid_request.n_x >= 1 and plan.grid_request.n_y >= 1 + + +def test_plan_none_when_no_loop(): + # crystal only, no loop -> no grid + assert _plan(_fake_mlbox(crystals=[(300, 300, 350, 350)])) is None + + +def test_loop_face_padding_grows_grid(monkeypatch): + box = (150, 150, 300, 300) + monkeypatch.setattr(mlb, "cfg_get", lambda key, default=None: 0.0 if "padding" in key else default) + unpadded = _plan(_fake_mlbox(loop_face=box)).grid_request + monkeypatch.setattr(mlb, "cfg_get", lambda key, default=None: 0.5 if "padding" in key else default) + padded = _plan(_fake_mlbox(loop_face=box)).grid_request + assert padded.n_x > unpadded.n_x and padded.n_y > unpadded.n_y + + +def test_crystal_union_extends_grid_only_when_enabled(monkeypatch): + # crystal extends well beyond the loop_face box on +x + mlbox = lambda: _fake_mlbox(loop_face=(150, 150, 300, 300), crystals=[(350, 150, 520, 300)]) + + monkeypatch.setattr(mlb, "cfg_get", lambda key, default=None: False if "include_crystal" in key else (0.0 if "padding" in key else default)) + off = _plan(mlbox()).grid_request + + monkeypatch.setattr(mlb, "cfg_get", lambda key, default=None: True if "include_crystal" in key else (0.0 if "padding" in key else default)) + on = _plan(mlbox()).grid_request + + assert on.n_x > off.n_x # grid widened to reach the crystal + assert on.n_y == off.n_y # crystal is within the loop's y-range + + +def test_zoom_box_uses_loop_all_unless_clipped(): + from aare.daq.operations.raster.service import RasterService + + svc = RasterService.__new__(RasterService) + + # loop_all fully inside the frame -> used for zoom + ok = mlb.MLRasterPlan( + grid_request=None, loop_all_box=(100, 100, 400, 400), + loop_face_box=(150, 150, 300, 300), image_width=1000, image_height=1000, + ) + assert svc._zoom_to_fit_box(ok) == (100, 100, 400, 400) + + # loop_all touches the left edge (clipped) -> fall back to loop_face + clipped = mlb.MLRasterPlan( + grid_request=None, loop_all_box=(0, 100, 400, 400), + loop_face_box=(150, 150, 300, 300), image_width=1000, image_height=1000, + ) + assert svc._zoom_to_fit_box(clipped) == (150, 150, 300, 300) + + # no loop_all -> loop_face + only_face = mlb.MLRasterPlan( + grid_request=None, loop_all_box=None, + loop_face_box=(150, 150, 300, 300), image_width=1000, image_height=1000, + ) + assert svc._zoom_to_fit_box(only_face) == (150, 150, 300, 300)