DAQ - MLBOX: adapting logic in MLBOX predictions now looks at multiple classes: face, crystal loop and pin.

daq/src/aaredaq/mlbox.py

@@ -1,25 +1,78 @@
 from pathlib import Path
-from typing import Tuple
+from typing import Tuple, Optional, List
 
 from ultralytics import YOLO
 
 
 class MlBox:
-    def __init__(self, model_path: str = "best.pt"):
+    def __init__(self, model_path: str = "new_best.pt"):
         model_path = str(Path(__file__).parent / model_path)
         self.__model = YOLO(model_path)
 
+        self.class_info = [
+            ["loop_all", (255, 0, 0)],    # class 0: Blue for loop_all
+            ["pin", (0, 255, 0)],         # class 1: Green for pin
+            ["crystal", (0, 0, 255)],     # class 2: Red for crystal
+            ["loop_face", (255, 255, 0)]  # class 3: Yellow for loop_face
+        ]
+
+
+
     @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) == 0:
-                    return x1, y1, x2, y2
-        return None
+                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:
+            return None
+        # if pin:
+        #     x1 = result[0]
+        #     pin_x1 = pin[0][0]
+        #     if result[0] > pin_x1:
+        #         return None
+        return result
 
-    def predict(self, image) -> None | Tuple[float, float, float, float]:
+    @staticmethod
+    def get_all_detections(results) -> List[Tuple[float, float, float, float, float, int]]:
+
+        all_detections = []
+
+        for r in results:
+            input_height, input_width = r.orig_shape
+            print(f"Model input size: {input_width}x{input_height}")
+            for box in r.boxes.data.tolist():
+                x1, y1, x2, y2, conf, cls = box
+                all_detections.append((x1, y1, x2, y2, conf, int(cls)))
+                print(f"Detection: class={int(cls)}, conf={conf:.2f}, box=({x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f})")
+
+        print(f"Total detections found: {len(all_detections)}")
+        return all_detections
+
+
+    def predict(self, image, filename: str | None = None) -> None | Tuple[float, float, float, float]:
+        print("running ml_box")
+        print("results:")
         results = self.__model.predict(source=image, conf=0.5)
-        pred = self.__get_pred(results)
 
+        detections = self.get_all_detections(results)
+
+        pred = self.__get_pred(results)
         return pred