mx/AareDAQ
0
1
Fork 0
Code Issues 49 Pull Requests Actions Packages Projects Releases 84 Wiki Activity

DAQ: the mlbox prediction bug fixes and updates, now uses target from aarelc_infer.py

Browse Source
This commit is contained in:
appleb_m
2026-04-17 17:27:04 +02:00
parent f4b5d420cf
commit e5194107d0
3 changed files with 292 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aare/common/aarelc_infer.py
+5 -3
View File
@@ -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}")
src/aare/daq/daq.py
+144 -29
View File
@@ -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:
src/aare/daq/mlbox.py
+143 -18
View File
@@ -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 {}