import argparse
import chess
import chess.pgn
from board_detector import find_board
from board_detector import find_pieces
from board_detector import init_global
from board_detector import display_img
from chess_ai import *
import move_translator
import cv2
import os
import time
# was having issues installing picamera2 on my PC
# if the raspi runs this code (with picamera2 installed), skip_camera will be false
# otherwise skip_camera is true and any camera code will be skipped
skip_camera = False
try:
from picamera2 import Picamera2, Preview
except ImportError:
skip_camera = True
class ChessGame:
def __init__(self, difficulty, color_scheme, show_cv, show_cam, loop, img_idx = 1, test_img = None, save_img_as = None):
self.board = chess.Board()
self.prev_board = chess.Board()
self.difficulty = difficulty
self.color_scheme = color_scheme
self.show_cv = show_cv
self.show_cam = show_cam
self.test_img = test_img
self.loop = loop
self.save_img_as = save_img_as
if img_idx:
self.img_idx = int(img_idx)
else:
self.img_idx = 1
# hard-coded borders to crop image
self.left_cut = 0
self.right_cut = 280
self.top_cut = 0
self.bottom_cut = 280
self.img_size = 512
if not skip_camera:
self.picam2 = Picamera2()
else:
self.picam2 = None
# self.picam2 = Picamera2()
def start_game(self):
print(f"Starting chess game (difficulty: {self.difficulty})")
# TODO - call initialize board in board_detector, initialize colors for color analysis,
# then loop until checkmate. also handle illegal moves (writing to screen if we end up doing that or just LEDs)
init_global(self.show_cv, skip_camera, self.img_size)
if self.save_img_as:
self.loop = True
# init camera
if not skip_camera:
preview_config = self.picam2.create_preview_configuration(main={"size": (2464, 2464)})
self.picam2.configure(preview_config)
if (self.show_cam):
self.picam2.start_preview(Preview.QTGL)
self.picam2.start()
elif not self.test_img:
self.test_img = 'ocr_ps_pink_yellow.jpg'
# initial setup of board
self.do_cv()
write_to("saved_FEN", self.board.fen())
while(1): # game loop
self.player_turn()
if self.check_game_over():
break
self.ai_turn()
if self.check_game_over():
break
# game is over
def player_turn(self):
# TODO - wait for user button, then check for valid move. loop until a valid move has been made
input("[Waiting to submit move. Replace with physical button]")
self.do_cv()
# handle cheating
while cheat_check(self.board.fen()):
print("CHECKING: ", self.board.fen())
self.board = self.prev_board.copy()
print("Player is cheating. Return to this state and try again:")
print(self.board)
input("[Waiting to submit move. Replace with physical button]")
self.do_cv()
def ai_turn(self):
info = chess_AI(self.board.fen(), 0)
# convert best move to coordinates to send
input("[SEND BEST MOVE TO ESP32 AND WAIT]")
self.do_cv()
def do_cv(self):
while(1): # essentially do_while(self.loop)
if (self.test_img):
img_txt = self.test_img + str(self.img_idx) + '.jpg'
img_path = os.path.join('ocr_test_images', img_txt)
orig_img = cv2.imread(img_path)
self.img_idx += 1
else:
orig_img = self.take_pic()
h,w,c = orig_img.shape
cropped_img = orig_img[self.top_cut:h-self.bottom_cut, self.left_cut:w-self.right_cut]
# cropped_img = orig_img
img = cv2.resize(cropped_img, (self.img_size, self.img_size))
# img = orig_img
if (self.show_cv and not self.loop):
# display_img([orig_img, img])
display_img([img])
if (self.loop):
answer = 0
while(answer != "y" and answer != "n"):
answer = input("Done looping? (y/n): ")
if answer == "y":
self.loop = False
elif answer == "n":
self.loop = True
if (not self.loop):
break
# warp the image based on the lines of the board
warped_img, sorted_warped_points = find_board(img)
if (warped_img is None):
return
# get the pieces based on color thresholding and easyocr
color_grid = find_pieces(warped_img, sorted_warped_points)
# convert color_grid to board, which we can get fen string from
if self.color_scheme == 'p/y':
self.color_grid_to_fen(color_grid, 'yellow', 'pink')
else:
self.color_grid_to_fen(color_grid, 'red', 'teal')
# check for invalid fen string
if not FEN_check(self.board.fen()):
print("Bad FEN.")
return
# print the board to terminal
print(self.board)
print(self.board.fen())
def check_game_over(self):
if self.board.is_checkmate():
print("Checkmate!")
return True
elif self.board.is_stalemate():
print("Stalemate.")
return True
elif self.board.is_insufficient_material():
print("Draw by insufficient material.")
return True
return False
def take_pic(self):
time.sleep(2) # camera needs this apparently
# name the image
if (self.save_img_as):
img_txt = self.save_img_as + str(self.img_idx) + '.jpg'
else:
img_txt = 'board' + str(self.img_idx) + '.jpg'
# save image
img_path = os.path.join('game_images', img_txt)
metadata = self.picam2.capture_file(img_path)
self.img_idx += 1
return cv2.imread(img_path)
def color_grid_to_fen(self, color_grid, color1, color2):
self.prev_board = self.board.copy()
for i, row in enumerate(color_grid):
for j, (color, _, _, letter) in enumerate(row):
piece_type = None
if letter is not None:
letter = letter[0]
if letter == "P":
piece_type = chess.PAWN
elif letter == "N":
piece_type = chess.KNIGHT
elif letter == "B":
piece_type = chess.BISHOP
elif letter == "R":
piece_type = chess.ROOK
elif letter == "Q":
piece_type = chess.QUEEN
elif letter == "K":
piece_type = chess.KING
piece_color = None
if color == color1:
piece_color = chess.BLACK
elif color == color2:
piece_color = chess.WHITE
if piece_color is not None and piece_type is not None:
self.board.set_piece_at(chess.square(j, 7 - i), chess.Piece(piece_type, piece_color))
continue
self.board.remove_piece_at(chess.square(j, 7 - i))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="AI Chess Robot with Computer Vision")
parser.add_argument("--difficulty", choices=["easy", "medium", "hard"],
default="medium", help="Chess AI difficulty (how far it looks ahead)")
parser.add_argument("--color_scheme", choices=["r/t", "p/y"],
default="p/y", help="Red and teal or pink and yellow for chess piece colors")
parser.add_argument("--show_cv", action="store_true", help="Show opencv images as processing occurs during game")
parser.add_argument("--show_cam", action="store_true", help="Show persistent camera view")
parser.add_argument("--loop", action="store_true", help="Loop before cv (for taking test images)")
parser.add_argument("--img_idx", help="Where to start indexing images (for naming them; default is 1 if not specified)")
parser.add_argument("--test_img", help="If specified, will use said image in test_images folder rather than camera input")
parser.add_argument("--save_img_as", help="If specified, will save image as given name in game_images")
args = parser.parse_args()
game = ChessGame(args.difficulty, args.color_scheme, args.show_cv, args.show_cam, args.loop, args.img_idx, args.test_img, args.save_img_as)
game.start_game()