diff --git a/daq/src/aaredaq/mlbox.py b/daq/src/aaredaq/mlbox.py index 369d62a9..7451307c 100644 --- a/daq/src/aaredaq/mlbox.py +++ b/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 \ No newline at end of file