Idk
0
rpi/board-detector/__init__.py
Normal file
31
rpi/board-detector/debug.py
Normal 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()
|
||||
79
rpi/board-detector/main.py
Normal 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()
|
||||
3
rpi/board-detector/paths.py
Normal 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"
|
||||
BIN
rpi/board-detector/test/1.png
Normal file
|
After Width: | Height: | Size: 111 KiB |
BIN
rpi/board-detector/test/2.webp
Normal file
|
After Width: | Height: | Size: 66 KiB |
BIN
rpi/board-detector/test/3.jpg
Normal file
|
After Width: | Height: | Size: 260 KiB |
BIN
rpi/board-detector/test/4.jpg
Normal file
|
After Width: | Height: | Size: 143 KiB |
BIN
rpi/board-detector/test/IMG_20251218_155300.jpg
Normal file
|
After Width: | Height: | Size: 160 KiB |
BIN
rpi/board-detector/test/glass-board.jpg
Normal file
|
After Width: | Height: | Size: 305 KiB |
BIN
rpi/board-detector/test/random_pieces.jpg
Normal file
|
After Width: | Height: | Size: 336 KiB |
BIN
rpi/board-detector/test/test.jpg
Normal file
|
After Width: | Height: | Size: 316 KiB |