Updates

rpi/board-detector/main.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
-from paths import *
+import os
+import random
 
 from ultralytics import YOLO
 import cv2
@@ -62,18 +63,31 @@ def prediction_to_fen(results, width, height):
 
 
 if __name__ == "__main__":
+    model_path = "../assets/models/unified-nano-refined.pt"
+    img_folder = "../training/datasets/pieces/unified/test/images/"
+    save_folder = "./results"
+    os.makedirs(save_folder, exist_ok=True)
 
-    img = cv2.imread(img_path)
-    height, width = img.shape[:2]
+    test_images = os.listdir(img_folder)
 
-    model = YOLO(model_path)
-    results = model.predict(source=img_path, conf=0.5)
+    for i in range(0, 10):
+        rnd = random.randint(0, len(test_images) - 1)
+        img_path = os.path.join(img_folder, test_images[rnd])
+        save_path = os.path.join(save_folder, test_images[rnd])
 
-    #fen = prediction_to_fen(results, height, width)
-    #print("Predicted FEN:", fen)
+        img = cv2.imread(img_path)
+        height, width = img.shape[:2]
 
-    annotated_image = results[0].plot()  # Annotated image as NumPy array
-    cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)  # make window resizable
-    cv2.imshow("YOLO Predictions", annotated_image)
-    cv2.waitKey(0)
-    cv2.destroyAllWindows()
+        model = YOLO(model_path)
+        results = model.predict(source=img_path, conf=0.5)
+
+        #fen = prediction_to_fen(results, height, width)
+        #print("Predicted FEN:", fen)
+
+        annotated_image = results[0].plot()
+        cv2.imwrite(save_path, annotated_image)
+        #cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
+        #cv2.imshow("YOLO Predictions", annotated_image)
+
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()

rpi/board-detector/paths.py

@@ -1,3 +0,0 @@
-model_path = "C:/Users/Laurent/Desktop/board-mate/rpi/assets/models/epoch-200.pt"
-#img_path = "./test/4.jpg"
-img_path = "../training/datasets/unified/train/images/WIN_20221220_11_27_27_Pro_jpg.rf.4f01cb68c8944ef1c4c7dc57847b4cd3.jpg"

rpi/board-detector/realtime_detect.py

@@ -1,11 +1,14 @@
 from ultralytics import YOLO
-from paths import *  # make sure model_path is defined here
 import cv2
 
 if __name__ == "__main__":
 
+    corner_model_path = "../assets/models/corner.pt"
+    pieces_model_path = "../assets/models/unified-nano-refined.pt"
+
     print("Initializing model...")
-    model = YOLO(model_path)
+    corner_model = YOLO(corner_model_path)
+    pieces_model = YOLO(pieces_model_path)
 
     print("Initializing camera...")
     cap = cv2.VideoCapture(0)
@@ -29,11 +32,13 @@ if __name__ == "__main__":
         # Optional: resize frame to improve YOLO performance
         # frame = cv2.resize(frame, (416, 416))
 
-        results = model.predict(source=frame, conf=0.5)
+        corner_result = corner_model.predict(source=frame, conf=0.6)
+        pieces_result = pieces_model.predict(source=frame, conf=0.6)
 
-        annotated_frame = results[0].plot()  # annotated frame as NumPy array
+        corner_annotated_frame = corner_result[0].plot()
+        pieces_annotated_frame = pieces_result[0].plot(img=corner_annotated_frame)
 
-        cv2.imshow("Predictions", annotated_frame)
+        cv2.imshow("Predictions", pieces_annotated_frame)
         cv2.resizeWindow("Predictions", 640, 640)
         if cv2.waitKey(1) & 0xFF == ord('q'):
             break