DAQ: update mlbox to get ml predictions via http

daq/src/aaredaq/mlbox.py

@@ -1,14 +1,17 @@
 from pathlib import Path
 from typing import Tuple, Optional, List
 
+import cv2
+import requests
 from ultralytics import YOLO
 
 
 class MlBox:
-    def __init__(self, model_path: str = "best_v8_18092025.pt"):
-        model_path = str(Path(__file__).parent / model_path)
-        self.__model = YOLO(model_path)
-
+    # # def __init__(self, model_path: str = "best_v8_18092025.pt"):
+    #     model_path = str(Path(__file__).parent / model_path)
+    #     self.__model = YOLO(model_path)
+    def __init__(self, url="http://mx-aare-test.psi.ch:8002/predict/?model=best_v8_18092025.pt"):
+        self.__url = url
         self.class_info = [
             ["loop_all", (255, 0, 0)],    # class 0: Blue for loop_all
             ["pin", (0, 255, 0)],         # class 1: Green for pin
@@ -16,39 +19,38 @@ class MlBox:
             ["loop_face", (255, 255, 0)]  # class 3: Yellow for loop_face
         ]
 
-
+    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")}
+        response = requests.post(self.__url, files=files)
+        return response.json()
 
     @staticmethod
-    def __get_pred(results) -> None | Tuple[float, float, float, float]:
-        loop_face = []
-        loop_all = []
-        pin = []
-        crystal = []
-        for r in results:
-            for box in r.boxes.data.tolist():
-                x1, y1, x2, y2, conf, cls = box
-                if int(cls) == 3:
-                    loop_face.append((x1, y1, x2, y2))
-                elif int(cls) == 2:
-                    crystal.append((x1, y1, x2, y2))
-                elif int(cls) == 1:
-                    pin.append((x1, y1, x2, y2))
-                elif int(cls) == 0:
-                    loop_all.append((x1, y1, x2, y2))
-        if crystal:
-            result = crystal[0]
-        elif loop_face:
-            result = loop_face[0]
-        elif loop_all:
-            result = loop_all[0]
-        else:
+    def __get_pred(results) -> None | Tuple[int, float, float, float, float]:
+        best_by_class: dict[int, tuple[float, float, float, float, float]] = {}
+
+        for pred in results:
+            cls = int(pred.get("class"))
+            conf = float(pred.get("confidence"))
+            box = pred.get("box") or {}
+            x1 = float(box.get("x1"))
+            y1 = float(box.get("y1"))
+            x2 = float(box.get("x2"))
+            y2 = float(box.get("y2"))
+
+            prev = best_by_class.get(cls)
+            if prev is None or conf > prev[4]:
+                best_by_class[cls] = (x1, y1, x2, y2, conf)
+
+        if not best_by_class:
             return None
-        # if pin:
-        #     x1 = result[0]
-        #     pin_x1 = pin[0][0]
-        #     if result[0] > pin_x1:
-        #         return None
-        return result
+
+        for preferred_cls in (2, 3, 0, 1):
+            if preferred_cls in best_by_class:
+                x1, y1, x2, y2, _ = best_by_class[preferred_cls]
+                return preferred_cls, x1, y1, x2, y2
 
     @staticmethod
     def get_all_detections(results) -> List[Tuple[float, float, float, float, float, int]]:
@@ -67,12 +69,15 @@ class MlBox:
         return all_detections
 
 
-    def predict(self, image, filename: str | None = None) -> None | Tuple[float, float, float, float]:
+    def predict(self, image, filename: str | None = None) -> None | Tuple[int, float, float, float, float]:
         print("running ml_box")
         print("results:")
-        results = self.__model.predict(source=image, conf=0.5)
-
-        detections = self.get_all_detections(results)
-
-        pred = self.__get_pred(results)
+        results = self.get_response(image)
+        print(results)
+        #results = self.__model.predict(source=image, conf=0.5)
+        #detections = self.get_all_detections(results.results)
+        preds = results.get("results") if isinstance(results, dict) else None
+        if not preds:
+            return None
+        pred = self.__get_pred(preds)
         return pred