mlbox debugging

src/aare/daq/mlbox.py

@@ -21,6 +21,8 @@ class BoxClassEnum(Enum):
     Pin = 1
     Crystal = 2
     Loop_face = 3
+    Ice = 4
+    Needle = 5
 
 class MlBox:
 
@@ -33,15 +35,17 @@ class MlBox:
             self.__url = "http://x10sa-spark-01.psi.ch:8002/predict/?model=best_yolo26l-seg-overlap-false_2026-03-16.engine"#v12_22092025.engine"
         elif bl == MXBeamline.X06SA:
             self.__url = ""
-            raise NotImplemented(f"MLBox not implemente for {bl}")
+            raise NotImplemented(f"MLBox not implemented for {bl}")
         else:
             raise Exception(f"unknown beamline {bl}")
 
     def get_response(self, image):
-        ok, buf = cv2.imencode(".jpg", image)
-        if not ok:
-            raise RuntimeError("Failed to encode image")
-        files={"file": ("image.jpg", buf.tobytes(), "image/jpeg")}
+        # cv2.imwrite("/sls/mx/applications/logs/image.png", image)
+        # ok, buf = cv2.imencode(".png", image)
+        # if not ok:
+        #     raise RuntimeError("Failed to encode image")
+        # files={"file": ("image.png", buf.tobytes(), "image/png")}
+        files = {"file": ("image.raw", image.tobytes(), "application/octet-stream")}
         response = requests.post(self.__url, files=files, timeout=10)
         response.raise_for_status()
         logger.debug(response.text)
@@ -162,9 +166,11 @@ class MlBox:
 
             current = out.boxes[base_key]
             if (current.conf or 0.0) < conf:
-                out.boxes[base_key] = MLBoxModel.from_tuple(cls, (x1, y1, x2, y2), conf)
-            else:
-                out.add_box(cls, (x1, y1, x2, y2), conf)
+                out.boxes[base_key] = MLBoxModel(
+                    cls=cls,
+                    box=BoundingBoxModel(top_x=x1, top_y=y1, bottom_x=x2, bottom_y=y2),
+                    conf=conf,
+                )
 
         return out if out.boxes else None
 
@@ -174,6 +180,67 @@ class MlBox:
             return None
         return best_by_class
 
+    @staticmethod
+    def get_preferred_class_box_with_confidence_threshold(
+            boxes: MLOutputModel,
+            preferred_class: Optional[Iterable[int] | int | MLBoxType] = None,
+            loop_preference_margin: float = 0.1
+    ) -> Optional[MLBoxModel]:
+        """
+        Get best box, preferring loops over pin even if pin has higher confidence,
+        unless pin's confidence exceeds loops by the margin.
+
+        Args:
+            boxes: MLOutputModel with detections
+            preferred_class: Override preference order (default: Crystal, Loop_face, Loop_all, Pin)
+            loop_preference_margin: Minimum confidence advantage pin needs to override loop preference (default 0.1)
+
+        Returns:
+            Best MLBoxModel according to preferences
+        """
+        if preferred_class is None:
+            order = (MLBoxType.Crystal, MLBoxType.Loop_face, MLBoxType.Loop_all, MLBoxType.Pin)
+        else:
+            if isinstance(preferred_class, MLBoxType):
+                order = (preferred_class,)
+            elif isinstance(preferred_class, int):
+                order = (MLBoxType(preferred_class),)
+            else:
+                order = tuple(MLBoxType(c) if isinstance(c, int) else c for c in preferred_class)
+
+        # Get best box from each class
+        best_boxes = {}
+        for cls in order:
+            m = boxes.get_best_for_class(cls)
+            if m:
+                best_boxes[cls] = m
+
+        if not best_boxes:
+            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)
+
+        best_loop = loop_face_box if loop_face_box else loop_all_box
+
+        if best_loop and pin_box:
+            pin_conf = pin_box.conf or 0.0
+            loop_conf = best_loop.conf or 0.0
+
+            # Only use pin if its confidence exceeds loop by margin
+            if pin_conf > (loop_conf + loop_preference_margin):
+                return pin_box
+            return best_loop
+
+        # Normal priority if no special case
+        for cls in order:
+            if cls in best_boxes:
+                return best_boxes[cls]
+
+        return None
+
     @staticmethod
     def get_preferred_class_box(boxes: MLOutputModel,
                                 preferred_class: Optional[Iterable[int] | int | MLBoxType] = None) -> Optional[MLBoxModel]: