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()