Idk

rpi/board-detector/main.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+from paths import *
+
+from ultralytics import YOLO
+import cv2
+
+# Map class names to FEN characters
+class_to_fen = {
+    'w_pawn': 'P',
+    'w_knight': 'N',
+    'w_bishop': 'B',
+    'w_rook': 'R',
+    'w_queen': 'Q',
+    'w_king': 'K',
+    'b_pawn': 'p',
+    'b_knight': 'n',
+    'b_bishop': 'b',
+    'b_rook': 'r',
+    'b_queen': 'q',
+    'b_king': 'k',
+}
+
+def prediction_to_fen(results, width, height):
+
+    # Initialize empty board
+    board = [['' for _ in range(8)] for _ in range(8)]
+
+    # Iterate through predictions
+    for result in results:
+        for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
+            x1, y1, x2, y2 = box.tolist()
+            class_name = model.names[int(cls)]
+            fen_char = class_to_fen.get(class_name)
+
+            if fen_char:
+                # Compute board square
+                col = int((x1 + x2) / 2 / (width / 8))
+                row = 7 - int((y1 + y2) / 2 / (height / 8))
+                board[row][col] = fen_char
+                print(f"[{class_name}] {fen_char} {row} {col}")
+
+    # Convert board to FEN
+    fen_rows = []
+    for row in board:
+        fen_row = ''
+        empty_count = 0
+        for square in row:
+            if square == '':
+                empty_count += 1
+            else:
+                if empty_count > 0:
+                    fen_row += str(empty_count)
+                    empty_count = 0
+                fen_row += square
+        if empty_count > 0:
+            fen_row += str(empty_count)
+        fen_rows.append(fen_row)
+
+    # Join rows into a FEN string (default: white to move, all castling rights, no en passant)
+    fen_string = '/'.join(fen_rows) + ' w KQkq - 0 1'
+    return fen_string
+
+
+if __name__ == "__main__":
+
+    img = cv2.imread(img_path)
+    height, width = img.shape[:2]
+
+    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()  # 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()