This commit is contained in:
2025-12-19 06:37:19 +01:00
parent 58e19bca64
commit 9c217c0599
29 changed files with 346 additions and 2 deletions

View File

View File

@@ -0,0 +1,31 @@
#!/usr/bin/env python3
import cv2
import numpy as np
from ultralytics import YOLO
from paths import *
def main():
model = YOLO(model_path)
# Load image
image = cv2.imread(img_path)
if image is None:
print(f"Failed to read {img_path}")
return
height, width = image.shape[:2]
warped = image # For now assume top-down view
# Run YOLO detection
results = model(warped)
res = results[0]
debug_img = res.plot() # draws boxes around detected objects
cv2.imshow("Detections", debug_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

View File

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

View File

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

Binary file not shown.

After

Width:  |  Height:  |  Size: 111 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 66 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 260 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 143 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 160 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 305 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 336 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 316 KiB