Minor fixes

rpi/models/detection/detector.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+import os
+import random
+
+import numpy as np
+from ultralytics import YOLO
+import cv2
+from ultralytics.engine.results import Results
+
+
+class Detector :
+
+    def __init__(self, model_path):
+        self.model = YOLO(model_path)
+        self.used_height = 640
+        self.used_width = 640
+
+    def make_prediction(self, image : str | np.ndarray) -> list[Results]:
+        return self.model.predict(source=image, conf=0.6)
+
+if __name__ == "__main__":
+    corner_model_path = "../../assets/models/edges.pt"
+    pieces_model_path = "../../assets/models/unified-nano-refined.pt"
+
+    corner_model = YOLO(corner_model_path)
+    pieces_model = YOLO(pieces_model_path)
+
+    img_folder = "../training/datasets/pieces/unified/test/images/"
+    save_folder = "./results"
+    os.makedirs(save_folder, exist_ok=True)
+
+    test_images = os.listdir(img_folder)
+
+    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])
+
+        image = cv2.imread(img_path)
+        height, width = image.shape[:2]
+
+        #fen = prediction_to_fen(results, height, width)
+        #print("Predicted FEN:", fen)
+
+        corner_result = corner_model.predict(source=image, conf=0.6)
+        pieces_result = pieces_model.predict(source=image, conf=0.6)
+
+        corner_annotated_image = corner_result[0].plot()
+        pieces_annotated_image = pieces_result[0].plot(img=corner_annotated_image)
+
+        cv2.imwrite(save_path, pieces_annotated_image)
+        #cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
+        #cv2.imshow("YOLO Predictions", annotated_image)
+
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()