DAQ: refactored MLBOX and loop_Centering MLBOXTYPe. Loop_centering should no longer move if the only class found is pin, ice or needle. Refactored tests.

src/aare/common/models.py

@@ -431,12 +431,12 @@ class FluorescenceSpectrumOutputModel(BaseModel):
 
 
 class MLBoxType(Enum):
-    Loop_all  = 0
-    Pin       = 1
-    Crystal   = 2
-    Loop_face = 3
-    Ice       = 4
-    Needle    = 5
+    LOOP_ALL  = 0
+    PIN       = 1
+    CRYSTAL   = 2
+    LOOP_FACE = 3
+    ICE       = 4
+    NEEDLE    = 5
 
 
 class BoundingBoxModel(BaseModel):
@@ -466,17 +466,17 @@ class MLOutputModel(BaseModel):
 
     @staticmethod
     def get_class_str(cls: MLBoxType) -> str:
-        if cls == MLBoxType.Loop_all:
+        if cls == MLBoxType.LOOP_ALL:
             return "Loop_all"
-        if cls == MLBoxType.Pin:
+        if cls == MLBoxType.PIN:
             return "Pin"
-        if cls == MLBoxType.Crystal:
+        if cls == MLBoxType.CRYSTAL:
             return "Crystal"
-        if cls == MLBoxType.Loop_face:
+        if cls == MLBoxType.LOOP_FACE:
             return "Loop_face"
-        if cls == MLBoxType.Ice:
+        if cls == MLBoxType.ICE:
             return "Ice"
-        if cls == MLBoxType.Needle:
+        if cls == MLBoxType.NEEDLE:
             return "Needle"
         return "Unknown"
 

src/aare/daq/daq.py

@@ -43,7 +43,7 @@ from aare.common.models import (
     SampleShortInfoList, AutofocusSettings,
     DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, ZoomModeEnum,
     SimpleScanParameters, MLBoxModel, FluorescenceSpectrumParameterModel,
-    FluorescenceSpectrumOutputModel, DAQOperation, LoopCenteringResult)
+    FluorescenceSpectrumOutputModel, DAQOperation, LoopCenteringResult, MLBoxType)
 from aare.common.automation_models import (
     AutomationProgress,
     StepState,
@@ -2362,7 +2362,7 @@ class AareDAQ:
         )
 
         return top_left, n_y
-
+    #TODO move wait_screenshot_sleep to save_screenshot_db!!!
     @log_timing(logger, "Loop center sequence")
     def __loop_center_sequence(
             self,
@@ -2441,10 +2441,11 @@ class AareDAQ:
                                 boxes=boxes,
                                 bundle_image=bundle_image,
                             )
-                            target = self._select_smargon_target(
-                                calculated_target=target,
-                                predicted_target=pred_target_point_smargon,
-                                tolerance_um=500.0)
+                            if cls in [MLBoxType.LOOP_ALL.value, MLBoxType.LOOP_FACE.value, MLBoxType.CRYSTAL.value]:
+                                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}")
@@ -2460,10 +2461,10 @@ class AareDAQ:
                                 found_classes_count[int(c)] = found_classes_count.get(int(c), 0) + 1
                         logger.debug(f"classes found: {classes}")
                         logger.debug(f"class found: {cls}")
-                        if cls is not None and cls != 1:
+                        if cls is not None and cls != MLBoxType.PIN.value and cls != MLBoxType.NEEDLE.value and cls != MLBoxType.ICE.value:
                             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)
@@ -2481,16 +2482,6 @@ class AareDAQ:
                                 )
                             self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{zoom_value:.0f}")
 
-                    # if targets_found_this_attempt == 0:
-                    #     logger.error(f"No targets found in this attempt {attempt}")
-                    #     raise LoopCenteringFailed
-                    #
-                    # else:
-                        # if targets_found_this_attempt >= len(base_angles):
-                        #     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}")
                     base_angles = (0, 90)
                     logger.debug(f"new base angles: {base_angles}")

src/aare/daq/mlbox.py

@@ -1,14 +1,11 @@
 from dataclasses import dataclass, field
-from typing import Optional, Iterable, TYPE_CHECKING, Any
+from typing import Optional, Iterable
 
 import cv2
 import numpy as np
 import time
 
-if TYPE_CHECKING:
-    from aarelcinfer_client.models import LatestPredictionModel
-else:
-    LatestPredictionModel = Any
+from aarelcinfer_client.models import LatestPredictionModel
 
 from aare.common.beamline import MXBeamline
 from aare.common.models import MLBoxModel, MLOutputModel, MLBoxType, BoundingBoxModel
@@ -266,7 +263,7 @@ class MlBox:
         return {"left": left, "top": top, "overlap_x": overlap_x, "overlap_y": overlap_y}
 
     def box_relation(self, boxes: MLOutputModel, classes:list[str]|None = None):
-        pin: MLBoxModel | None = boxes.get_best_for_class(MLBoxType.Pin)
+        pin: MLBoxModel | None = boxes.get_best_for_class(MLBoxType.PIN)
 
         if not pin or not classes:
             return None
@@ -305,7 +302,7 @@ class MlBox:
     def _filter_predictions(self, predictions: MLOutputModel, overlap_with_pin: Optional[float] = None,
                             confidence_min: Optional[float] = None):
 
-        pin = predictions.get_best_for_class(MLBoxType.Pin)
+        pin = predictions.get_best_for_class(MLBoxType.PIN)
         keys_to_remove = []
         for key, model in predictions.boxes.items():
             if model is None or model.box is None:
@@ -314,7 +311,7 @@ class MlBox:
             if confidence_min is not None and (model.conf or 0.0) < confidence_min:
                 keys_to_remove.append(key)
                 continue
-            if model.cls == MLBoxType.Pin:
+            if model.cls == MLBoxType.PIN:
                 continue
             if overlap_with_pin is not None and pin:
                 if self._check_overlap(model, pin) >= overlap_with_pin:
@@ -410,7 +407,7 @@ class MlBox:
             Best MLBoxModel according to preferences
         """
         if preferred_class is None:
-            order = (MLBoxType.Crystal, MLBoxType.Loop_face, MLBoxType.Loop_all, MLBoxType.Pin)
+            order = (MLBoxType.CRYSTAL, MLBoxType.LOOP_FACE, MLBoxType.LOOP_ALL, MLBoxType.PIN)
         else:
             if isinstance(preferred_class, MLBoxType):
                 order = (preferred_class,)
@@ -430,9 +427,9 @@ class MlBox:
             return None
 
         # Special handling: prefer loops over pin unless pin is significantly better
-        pin_box = best_boxes.get(MLBoxType.Pin)
-        loop_face_box = best_boxes.get(MLBoxType.Loop_face)
-        loop_all_box = best_boxes.get(MLBoxType.Loop_all)
+        pin_box = best_boxes.get(MLBoxType.PIN)
+        loop_face_box = best_boxes.get(MLBoxType.LOOP_FACE)
+        loop_all_box = best_boxes.get(MLBoxType.LOOP_ALL)
 
         best_loop = loop_face_box if loop_face_box else loop_all_box
 
@@ -457,7 +454,7 @@ class MlBox:
                                 preferred_class: Optional[Iterable[int] | int | MLBoxType] = None) -> Optional[
         MLBoxModel]:
         if preferred_class is None:
-            order = (MLBoxType.Crystal, MLBoxType.Loop_face, MLBoxType.Loop_all, MLBoxType.Pin)
+            order = (MLBoxType.CRYSTAL, MLBoxType.LOOP_FACE, MLBoxType.LOOP_ALL, MLBoxType.PIN)
         else:
             if isinstance(preferred_class, MLBoxType):
                 order = (preferred_class,)