diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py index 84d00fc7..18696653 100644 --- a/src/aare/daq/mlbox.py +++ b/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]: