From e5194107d0af512c446049d446f31d0bf8ef2d5b Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 17 Apr 2026 17:27:04 +0200 Subject: [PATCH] DAQ: the mlbox prediction bug fixes and updates, now uses target from aarelc_infer.py --- src/aare/common/aarelc_infer.py | 8 +- src/aare/daq/daq.py | 173 ++++++++++++++++++++++++++------ src/aare/daq/mlbox.py | 161 +++++++++++++++++++++++++---- 3 files changed, 292 insertions(+), 50 deletions(-) diff --git a/src/aare/common/aarelc_infer.py b/src/aare/common/aarelc_infer.py index 80df8d7d..d37e06ec 100644 --- a/src/aare/common/aarelc_infer.py +++ b/src/aare/common/aarelc_infer.py @@ -99,6 +99,8 @@ if __name__ == "__main__": img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) # Draw Overlay using the model attributes + print(prediction.target_point.x, prediction.target_point.y) + print("boxes: ", prediction.boxes) for det in prediction.boxes: x1, y1 = int(det.x1), int(det.y1) x2, y2 = int(det.x2), int(det.y2) @@ -118,9 +120,9 @@ if __name__ == "__main__": pts = pts.reshape((-1, 1, 2)) cv2.polylines(img, [pts], isClosed=True, color=(0, 0, 255), thickness=2) - cv2.imshow("Latest Prediction Bundle", img) - cv2.waitKey(0) - cv2.destroyAllWindows() + # cv2.imshow("Latest Prediction Bundle", img) + # cv2.waitKey(0) + # cv2.destroyAllWindows() except Exception as e: print(f"Error processing prediction bundle: {e}") \ No newline at end of file diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index c9731eb7..7cd5881e 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -20,7 +20,7 @@ from aare.common.autofocus_tools import focus_measure_edges from aare.daq.config import BeamlineConfig, ABR_POS_MOUNT, ABR_OMEGA_MOUNT from aare.daq.config import BeamlineStateEnum from aare.daq.devices import BeamlineDevices -from aare.daq.mlbox import MlBox +from aare.daq.mlbox import MlBox, MLBoxPredictionResult, MLBoxPredictionsResult from aare.common.beamline import MXBeamline from aare.common.coordinate import Coordinate, SmargonCoordinate, AerotechCoordinate from aare.common.diffraction_geometry import DiffractionGeometry @@ -82,6 +82,20 @@ class AareDAQ: self._face_detection_progress_cb(payload) except Exception as e: logger.warning(f"Failed to emit face detection progress: {e}") + + def _log_ml_bundle_meta( + self, + context: str, + *, + target_point: tuple[float, float] | None = None, + focus: float | None = None, + ) -> None: + if target_point is None and focus is None: + return + logger.debug( + f"ML bundle metadata for {context}: target_point={target_point}, focus={focus}" + ) + #TODO make sure this is implemented currectly def _handle_operation_error(self, operation_name: str, sample: SampleShortInfo | None, error: Exception, error_type: str = "generic") -> None: @@ -109,6 +123,7 @@ class AareDAQ: self.__aare.sample_failed(sample, failed_comment=f"Error in {operation_name}: {error}") except Exception as db_error: logger.error(f"Failed to report {operation_name} error to database: {db_error}") + #todo make sure these functions are correctly implemented! def _execute_mount_and_prepare(self, sample: SampleShortInfo) -> bool: """ @@ -1217,7 +1232,19 @@ class AareDAQ: def __ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None) -> RasterGridRequest | None: time.sleep(0.2) # Just to be sure image is stable - m, bundle_image = self.__mlbox.predict(preferred_class=(3, 0), return_image=True) + prediction_result: MLBoxPredictionResult = self.__mlbox.predict( + preferred_class=(3, 0), + return_image=True, + return_bundle_meta=True, + ) + m = prediction_result.box + bundle_image = prediction_result.image + self._log_ml_bundle_meta( + f"ml_bounding_box:{filename or 'unnamed'}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) + if m is None: if filename is not None and bundle_image is not None: self.__aare.upload_image(sample_id, f"{filename}_no_detection", bundle_image) @@ -1248,13 +1275,27 @@ class AareDAQ: omega_deg=geom.omega_deg ) - def __ml_loop_centre_box(self, sample_id: int | None = None, filename: str | None = None) -> tuple[ - SmargonCoordinate | None, int | None, list[int] | None]: - boxes, bundle_image = self.__mlbox.predict_all_best( - overlap_with_pin=0.5, - confidence_min=0.3, - return_image=True - ) + def __ml_loop_centre_box( + self, + sample_id: int | None = None, + filename: str | None = None, + boxes=None, + bundle_image: np.ndarray | None = None, + ) -> tuple[SmargonCoordinate | None, int | None, list[int] | None]: + if boxes is None: + prediction_result: MLBoxPredictionsResult = self.__mlbox.predict_all_best( + overlap_with_pin=0.5, + confidence_min=0.3, + return_image=True, + return_bundle_meta=True + ) + boxes = prediction_result.predictions + bundle_image = prediction_result.image + self._log_ml_bundle_meta( + f"ml_loop_centre_box:{filename or 'unnamed'}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) if boxes is None: if filename is not None and bundle_image is not None: @@ -1285,7 +1326,7 @@ class AareDAQ: geom = self.sample_geometry - if cls == 0: # loop_all + if cls == 0: # loop_all if y1 + y2 <= x1 + x2: centre_y = y1 + (y2 - y1) / 2 centre_x = x1 @@ -1303,13 +1344,13 @@ class AareDAQ: centre_y = y1 centre_x = x1 + (x2 - x1) / 2 - elif cls == 1: # pin - centre_y = y1 + (y2 - y1)/2 + elif cls == 1: # pin + centre_y = y1 + (y2 - y1) / 2 centre_x = x1 - elif cls == 2 or cls == 3: #crystal or loop_face - centre_y = y1 + (y2 - y1)/2 - centre_x = x1 + (x2 - x1)/2 + elif cls == 2 or cls == 3: # crystal or loop_face + centre_y = y1 + (y2 - y1) / 2 + centre_x = x1 + (x2 - x1) / 2 else: logger.debug(f"unknown box class {cls}") @@ -1412,7 +1453,17 @@ class AareDAQ: logger.info(f"time to rotate 15 degrees: {time.perf_counter() - rotate_time}") box_time = time.perf_counter() - m, bundle_image = self.__mlbox.predict(preferred_class=(3, 0), return_image=True) + prediction_result: MLBoxPredictionResult = self.__mlbox.predict( + preferred_class=(3, 0), + return_image=True, + return_bundle_meta=True + ) + m = prediction_result.box + self._log_ml_bundle_meta( + f"face_detection_angle_{angle}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) logger.info(f"time to predict: {time.perf_counter() - box_time}") if not m or not m.box: @@ -1506,6 +1557,43 @@ class AareDAQ: self._emit_face_detection_progress(result) return result + @staticmethod + def _select_smargon_target( + calculated_target: SmargonCoordinate | None, + predicted_target: SmargonCoordinate | None, + tolerance_um: float = 500.0, + ) -> SmargonCoordinate | None: + if calculated_target is None: + return predicted_target + if predicted_target is None: + return calculated_target + + tolerance_mm = tolerance_um / 1000.0 + dx = abs(calculated_target.sh_mm.x - predicted_target.sh_mm.x) + dy = abs(calculated_target.sh_mm.y - predicted_target.sh_mm.y) + dz = abs(calculated_target.sh_mm.z - predicted_target.sh_mm.z) + + if dx <= tolerance_mm and dy <= tolerance_mm and dz <= tolerance_mm: + logger.info( + "Using prediction target because it is within %.0f um of calculated target " + "(dx=%.4f mm, dy=%.4f mm, dz=%.4f mm)", + tolerance_um, + dx, + dy, + dz, + ) + return predicted_target + + logger.info( + "Keeping calculated target because prediction target is outside %.0f um " + "(dx=%.4f mm, dy=%.4f mm, dz=%.4f mm)", + tolerance_um, + dx, + dy, + dz, + ) + return calculated_target + 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 @@ -1513,24 +1601,17 @@ class AareDAQ: found_classes_count: dict[int, int] = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0} try: - self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter - # zoom_settings = self.__cfg.zoom_settings.z - + start = time.perf_counter() 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, 90) if (zoom_iter % 2 == 0) else (90, 0) - if sample_id is not None: + if sample_id is not None and zoom_iter == 0: 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") while attempt < max_attempt: - # self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain) - # exp = int(exposure * 1000) - # gn = int(gain) self.zoom = zoom_value found_flag = False @@ -1546,8 +1627,41 @@ class AareDAQ: filename = f"{sample_id}_{angle}_{zoom_value:.0f}" if sample_id is not None else None try: - self.save_screenshot(filename=f'{sample_id}_{angle}') - target, cls, classes = self.__ml_loop_centre_box(sample_id, filename) + time_to_get_pred = time.perf_counter() + prediction_result: MLBoxPredictionsResult = self.__mlbox.predict_all_best( + overlap_with_pin=0.5, + confidence_min=None, + return_image=True, + return_bundle_meta=True + ) + logger.info(f"time to move: {time.perf_counter() - time_to_get_pred}") + boxes = prediction_result.predictions + pred_target_point = prediction_result.target_point + + pred_target_point_smargon = None + if pred_target_point is not None: + pred_target_point_smargon = SmargonCoordinate(sh_mm=self.sample_geometry.picture_to_smargon( + Coordinate(x=pred_target_point[0], y=pred_target_point[1]) + )) + + bundle_image = prediction_result.image + self._log_ml_bundle_meta( + f"loop_center_angle_{angle}_zoom_{zoom_value:.0f}", + target_point=prediction_result.target_point, + focus=prediction_result.focus, + ) + + target, cls, classes = self.__ml_loop_centre_box( + sample_id=sample_id, + filename=filename, + boxes=boxes, + bundle_image=bundle_image, + ) + target = self._select_smargon_target( + calculated_target=target, + predicted_target=pred_target_point_smargon, + tolerance_um=500.0) + logger.debug(f"calculated target: {target} compares to prediction: {pred_target_point_smargon}") except Exception as e: logger.error(f"Error getting ML box for angle {angle}") logger.error(f"Exception: {e}") @@ -1566,7 +1680,6 @@ class AareDAQ: targets_found_this_attempt += 1 found_flag = True found_angle = angle - time_to_move_smargon = time.perf_counter() self.__devs.smargon_pos = target self.__devs.smargon_wait(60) @@ -1585,7 +1698,8 @@ class AareDAQ: else: if targets_found_this_attempt >= len(base_angles): - logger.debug(f"sucessfully found {targets_found_this_attempt} targets in attempt {attempt + 1} ") + logger.debug( + f"sucessfully found {targets_found_this_attempt} targets in attempt {attempt + 1} ") 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}") @@ -1605,6 +1719,7 @@ class AareDAQ: 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") + logger.debug(f"time to loop center: {time.perf_counter() - start}") return True except Exception as e: diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py index e8655b6d..9ba7f2b1 100644 --- a/src/aare/daq/mlbox.py +++ b/src/aare/daq/mlbox.py @@ -1,13 +1,10 @@ -import io -from enum import Enum +from dataclasses import dataclass, field + from typing import Optional, Iterable import cv2 import numpy as np import time -import requests - -from PIL import Image from aarelcinfer_client.models import LatestPredictionModel @@ -19,6 +16,35 @@ from aare.common.logger_config import setup_logger logger=setup_logger("aareDAQ") +@dataclass +class MLBundleMeta: + target_point: tuple[float, float] | None = None + focus: float | None = None + + +@dataclass +class MLBoxPredictionResult: + box: MLBoxModel | None + image: np.ndarray | None + target_point: tuple[float, float] | None = None + focus: float | None = None + + +@dataclass +class MLBoxPredictionsResult: + predictions: MLOutputModel | None + image: np.ndarray | None + target_point: tuple[float, float] | None = None + focus: float | None = None + + +@dataclass +class MLBundle: + predictions: MLOutputModel | None = None + image: np.ndarray | None = None + bundle_meta: MLBundleMeta = field(default_factory=MLBundleMeta) + + class MlBox: RETRY_COUNT = 3 @@ -88,20 +114,74 @@ class MlBox: time.sleep(self.RETRY_SLEEP_S) raise last_error - def _collect_best_bundle(self, attempts: int = RETRY_COUNT) -> tuple[MLOutputModel | None, np.ndarray | None]: + @staticmethod + def _extract_target_point(prediction: LatestPredictionModel | None) -> tuple[float, float] | None: + if prediction is None: + return None + + raw = getattr(prediction, "target_point", None) + if raw is None: + return None + + try: + if isinstance(raw, dict): + x = raw.get("x") + y = raw.get("y") + if x is not None and y is not None: + return float(x), float(y) + + if hasattr(raw, "x") and hasattr(raw, "y"): + return float(raw.x), float(raw.y) + + if isinstance(raw, (list, tuple)) and len(raw) >= 2: + return float(raw[0]), float(raw[1]) + except Exception as e: + logger.warning(f"Failed to parse target_point from prediction metadata: {e}") + + return None + + @staticmethod + def _extract_bundle_meta(prediction: LatestPredictionModel | None) -> MLBundleMeta: + target_point = MlBox._extract_target_point(prediction) + + focus = None + if prediction is not None: + try: + raw_focus = getattr(prediction, "focus_score", None) + if raw_focus is not None: + focus = float(raw_focus) + except Exception as e: + logger.warning(f"Failed to parse focus_score from prediction metadata: {e}") + + return MLBundleMeta(target_point=target_point, focus=focus) + + def _extract_bundle_candidate(self, bundle) -> MLBundle: + prediction = self._coerce_prediction_metadata(getattr(bundle, "metadata", None)) + image = self._decode_bundle_image(getattr(bundle, "image_jpeg", None)) + predictions = self._all_from_prediction_model(prediction) + bundle_meta = self._extract_bundle_meta(prediction) + return MLBundle(predictions=predictions, image=image, bundle_meta=bundle_meta) + + def _collect_best_bundle(self, attempts: int = RETRY_COUNT) -> MLBundle: best_predictions: MLOutputModel | None = None best_image: np.ndarray | None = None best_score = (0, 0.0) + best_meta = MLBundleMeta() for attempt in range(1, attempts + 1): bundle = self._fetch_prediction_bundle() - prediction = self._coerce_prediction_metadata(getattr(bundle, "metadata", None)) - image = self._decode_bundle_image(getattr(bundle, "image_jpeg", None)) - predictions = self._all_from_prediction_model(prediction) + ml_bundle = self._extract_bundle_candidate(bundle) + predictions, image, bundle_meta = ml_bundle.predictions, ml_bundle.image, ml_bundle.bundle_meta if best_image is None and image is not None: best_image = image + if best_meta.target_point is None and bundle_meta.target_point is not None: + best_meta.target_point = bundle_meta.target_point + + if best_meta.focus is None and bundle_meta.focus is not None: + best_meta.focus = bundle_meta.focus + score = self._prediction_score(predictions) logger.debug( f"Bundle candidate {attempt}/{attempts}: detections={score[0]}, max_conf={score[1]:.3f}" @@ -112,6 +192,7 @@ class MlBox: best_predictions = predictions if image is not None: best_image = image + best_meta = bundle_meta if score[0] > 0: logger.debug("Using non-empty prediction bundle candidate") @@ -125,7 +206,7 @@ class MlBox: elif not best_predictions.boxes: logger.info("Prediction bundle candidates contained no supported detections") - return best_predictions, best_image + return MLBundle(predictions=best_predictions, image=best_image, bundle_meta=best_meta) def _get_latest_bundle_image(self) -> np.ndarray | None: bundle = self._fetch_prediction_bundle() @@ -376,18 +457,45 @@ class MlBox: def predict(self, preferred_class = None, overlap_with_pin: float | None = None, confidence_min: float | None = None, - return_image: bool = False - ) -> None | MLBoxModel | tuple[MLBoxModel | None, np.ndarray | None]: - predictions, image = self._collect_best_bundle() + return_image: bool = False, return_bundle_meta: bool = False + ) -> None | MLBoxModel | tuple[MLBoxModel | None, np.ndarray | None] | MLBoxPredictionResult: + ml_bundle = self._collect_best_bundle() + predictions, image, bundle_meta = ml_bundle.predictions, ml_bundle.image, ml_bundle.bundle_meta if not predictions: + if return_bundle_meta: + return MLBoxPredictionResult( + box=None, + image=image if return_image else None, + target_point=bundle_meta.target_point, + focus=bundle_meta.focus, + ) return (None, image) if return_image else None self._filter_predictions(predictions=predictions, overlap_with_pin=overlap_with_pin, confidence_min=confidence_min) box = self.get_preferred_class_box(predictions, preferred_class) + if return_bundle_meta: + return MLBoxPredictionResult( + box=box, + image=image if return_image else None, + target_point=bundle_meta.target_point, + focus=bundle_meta.focus, + ) return (box, image) if return_image else box - def predict_best_no_filter(self, return_image: bool = False) -> Optional[MLOutputModel] | tuple[Optional[MLOutputModel], np.ndarray | None]: - predictions, image = self._collect_best_bundle() + def predict_best_no_filter( + self, + return_image: bool = False, + return_bundle_meta: bool = False + ) -> Optional[MLOutputModel] | tuple[Optional[MLOutputModel], np.ndarray | None] | MLBoxPredictionsResult: + ml_bundle = self._collect_best_bundle() + predictions, image, bundle_meta = ml_bundle.predictions, ml_bundle.image, ml_bundle.bundle_meta + if return_bundle_meta: + return MLBoxPredictionsResult( + predictions=predictions, + image=image if return_image else None, + target_point=bundle_meta.target_point, + focus=bundle_meta.focus, + ) if return_image: return predictions, image return predictions @@ -395,14 +503,30 @@ class MlBox: def predict_all_best(self, overlap_with_pin: float | None = None, confidence_min: float | None = None, - return_image: bool = False) -> Optional[MLOutputModel] | tuple[Optional[MLOutputModel], np.ndarray | None]: - best, image = self._collect_best_bundle() + return_image: bool = False, + return_bundle_meta: bool = False) -> Optional[MLOutputModel] | tuple[Optional[MLOutputModel], np.ndarray | None] | MLBoxPredictionsResult: + ml_bundle = self._collect_best_bundle() + best, image, bundle_meta = ml_bundle.predictions, ml_bundle.image, ml_bundle.bundle_meta if not best: logger.debug(f"No best predictions from ML bundle: {best}") + if return_bundle_meta: + return MLBoxPredictionsResult( + predictions=None, + image=image if return_image else None, + target_point=bundle_meta.target_point, + focus=bundle_meta.focus, + ) return (None, image) if return_image else None logger.debug(f"Best predictions from ML bundle: {best}") self._filter_predictions(best, overlap_with_pin=overlap_with_pin, confidence_min=confidence_min) logger.debug(f"Filtered best predictions from ML bundle: {best}") + if return_bundle_meta: + return MLBoxPredictionsResult( + predictions=best, + image=image if return_image else None, + target_point=bundle_meta.target_point, + focus=bundle_meta.focus, + ) return (best, image) if return_image else best def predict_all(self, @@ -414,7 +538,8 @@ class MlBox: The ordinal is the running count per class (1-based) in the order they appear after filtering. Example keys: 'Crystal_1', 'Crystal_2', 'Loop_face_1', 'Pin_1', ... """ - grouped, image = self._collect_best_bundle() + ml_bundle = self._collect_best_bundle() + grouped, image, _bundle_meta = ml_bundle.predictions, ml_bundle.image, ml_bundle.bundle_meta if not grouped: return ({}, image) if return_image else {}