From e7201f7c0fdfe5e378c06b3f2935dfdd2f5bc517 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 22 May 2026 15:23:45 +0200 Subject: [PATCH] DAQ: added common operation helpeers --- src/aare/daq/operations/common/__init__.py | 0 .../daq/operations/common/ml_bounding_box.py | 214 ++++++++++++++++++ .../operations/common/simulate_scan_result.py | 63 ++++++ 3 files changed, 277 insertions(+) create mode 100644 src/aare/daq/operations/common/__init__.py create mode 100644 src/aare/daq/operations/common/ml_bounding_box.py create mode 100644 src/aare/daq/operations/common/simulate_scan_result.py diff --git a/src/aare/daq/operations/common/__init__.py b/src/aare/daq/operations/common/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/aare/daq/operations/common/ml_bounding_box.py b/src/aare/daq/operations/common/ml_bounding_box.py new file mode 100644 index 00000000..19d4df69 --- /dev/null +++ b/src/aare/daq/operations/common/ml_bounding_box.py @@ -0,0 +1,214 @@ +import time +from math import ceil, floor +from typing import Callable + +import cv2 + +from aare.common.coordinate import Coordinate, SmargonCoordinate +from aare.common.exception_handler import AutoRasterSampleSkipped +from aare.common.logger_events import ( + geom_log_context, + log_ml_bundle_meta, + merge_log_context, + sample_log_context, +) +from aare.common.models import SampleShortInfo +from aare.common.raster_grid import RasterGridRequest +from aare.common.sample_geometry import SampleGeometryModel +from aare.daq.mlbox import MLBoxPredictionResult, MlBox + + +def scale_auto_raster_grid( + *, + n_x: int, + n_y: int, + grid_size: Coordinate, + max_images: int, + min_cell_size_mm: float, + skip: bool, +) -> tuple[int, int, Coordinate]: + n_x = max(1, int(n_x)) + n_y = max(1, int(n_y)) + + image_count = n_x * n_y + if image_count <= max_images: + return n_x, n_y, grid_size + + if skip: + raise AutoRasterSampleSkipped( + f"Auto-raster grid has {image_count} images, which exceeds " + f"the automation limit of {max_images}; skipping sample" + ) + + physical_size_x_mm = n_x * grid_size.x + physical_size_y_mm = n_y * grid_size.y + + scale = (image_count / max_images) ** 0.5 + scaled_n_x = max(1, int(floor(n_x / scale))) + scaled_n_y = max(1, int(floor(n_y / scale))) + + while scaled_n_x * scaled_n_y > max_images: + if scaled_n_x >= scaled_n_y and scaled_n_x > 1: + scaled_n_x -= 1 + elif scaled_n_y > 1: + scaled_n_y -= 1 + else: + break + + scaled_grid_size = Coordinate( + x=max(min_cell_size_mm, physical_size_x_mm / scaled_n_x), + y=max(min_cell_size_mm, physical_size_y_mm / scaled_n_y), + ) + return scaled_n_x, scaled_n_y, scaled_grid_size + + +def get_ml_bounding_box( + *, + 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, +) -> RasterGridRequest | None: + time.sleep(0.2) + + sample_id = getattr(sample, "db_id", None) + + prediction_result: MLBoxPredictionResult = mlbox.predict( + preferred_class=(3, 0), + return_image=True, + return_bundle_meta=True, + ) + m = prediction_result.box + bundle_image = prediction_result.image + + log_ml_bundle_meta( + logger, + f"ml_bounding_box:{filename or 'unnamed'}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) + + if m is None: + logger.warning( + "ML bounding box returned no detection", + extra={ + "sample_id": sample_id, + "ml_image_name": filename, + "target_point": prediction_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 = m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y + + 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) + + geom = sample_geometry + logger.info( + "ML bounding box selected", + extra=merge_log_context( + sample_log_context(sample), + { + "sample_id": sample_id, + "ml_image_name": filename, + }, + geom_log_context(geom), + { + "box_x1": x1, + "box_y1": y1, + "box_x2": x2, + "box_y2": y2, + "target_point": prediction_result.target_point, + }, + ), + ) + + 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))) + n_y = max(1, abs(ceil((y2 - y1) * geom.pixel_in_mm / grid_size.y))) + + original_n_x = n_x + original_n_y = n_y + original_grid_size = grid_size + + n_x, n_y, grid_size = scale_auto_raster_grid( + n_x=n_x, + n_y=n_y, + grid_size=grid_size, + max_images=max_images, + min_cell_size_mm=min_cell_size_mm, + skip=skip_if_exceed_max_image_threshold, + ) + + if (n_x, n_y, grid_size.x, grid_size.y) != ( + original_n_x, + original_n_y, + original_grid_size.x, + original_grid_size.y, + ): + logger.info( + "Scaled ML raster grid to stay within auto-raster image limit", + extra=merge_log_context( + sample_log_context(sample), + { + "sample_id": sample_id, + "ml_image_name": filename, + "max_images": max_images, + "original_n_x": original_n_x, + "original_n_y": original_n_y, + "original_image_count": original_n_x * original_n_y, + "original_grid_size_x_mm": original_grid_size.x, + "original_grid_size_y_mm": original_grid_size.y, + "scaled_n_x": n_x, + "scaled_n_y": n_y, + "scaled_image_count": n_x * n_y, + "scaled_grid_size_x_mm": grid_size.x, + "scaled_grid_size_y_mm": grid_size.y, + }, + ), + ) + + logger.info( + "Converted ML bounding box to raster request", + extra=merge_log_context( + sample_log_context(sample), + { + "sample_id": sample_id, + "ml_image_name": filename, + "start_sh_x_mm": start_coord.x, + "start_sh_y_mm": start_coord.y, + "start_sh_z_mm": start_coord.z, + "grid_size_x_mm": grid_size.x, + "grid_size_y_mm": grid_size.y, + "n_x": n_x, + "n_y": n_y, + "smargon_phi_deg": geom.smargon.phi_deg, + "smargon_chi_deg": geom.smargon.chi_deg, + }, + ), + ) + + return RasterGridRequest( + exp_time_s=0.01, + transmission=1.0, + smargon_top_left=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, + ) \ No newline at end of file diff --git a/src/aare/daq/operations/common/simulate_scan_result.py b/src/aare/daq/operations/common/simulate_scan_result.py new file mode 100644 index 00000000..cef36566 --- /dev/null +++ b/src/aare/daq/operations/common/simulate_scan_result.py @@ -0,0 +1,63 @@ +import copy + +from jfjoch_client import ScanResult, ScanResultImagesInner + +from aare.common.raster_grid import CompletedRasterGridElem, RasterGridRequest +from aare.common.rotation_scan import CompletedRotationScan, RotationScanRequest + + +def build_fake_scan_result( + *, + file_prefix: str | None, + image_count: int, + rotation: bool = False, + start_angle: float = 0.0, +) -> ScanResult: + total_images = max(1, image_count) + angle_step = 360.0 / total_images if rotation else 0.0 + + images = [ + ScanResultImagesInner( + number=i, + efficiency=1.0, + bkg=0.0, + spots=0, + spots_low_res=0, + spots_indexed=0, + index=0, + b=0.0, + angle=start_angle + i * angle_step if rotation else None, + ) + for i in range(total_images) + ] + return ScanResult(file_prefix=file_prefix, images=images) + + +def build_fake_rotation_result( + request: RotationScanRequest, + *, + start_angle: float = 0.0, +) -> CompletedRotationScan: + result = build_fake_scan_result( + file_prefix=request.file_prefix, + image_count=request.steps, + rotation=True, + start_angle=start_angle, + ) + return CompletedRotationScan( + request=copy.deepcopy(request), + result=result, + ) + + +def build_fake_raster_result(request: RasterGridRequest) -> CompletedRasterGridElem: + result = build_fake_scan_result( + file_prefix=request.file_prefix, + image_count=request.n_x * request.n_y, + rotation=False, + ) + return CompletedRasterGridElem( + request=copy.deepcopy(request), + result=result, + centre_of_mass=None, + )