game.py

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
import chess_ai
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, show_cv, show_cam, loop, img_idx = 1, test_img = None, save_img_as = None):
        self.board = chess.Board()
        self.difficulty = difficulty
        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 = 290
        self.top_cut = 0
        self.bottom_cut = 290

        self.img_size = 512

        if not skip_camera:
            self.picam2 = Picamera2()
        else:
            self.picam2 = None

    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

        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_smallp2.jpg'

        while(1): # essentially do_while(self.loop)
            if (self.test_img):
                img_path = os.path.join('ocr_test_images', self.test_img)
                orig_img = cv2.imread(img_path)
            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

        warped_img, sorted_warped_points = find_board(img)

        if (warped_img is None):
            return
        
        color_grid = find_pieces(warped_img, sorted_warped_points)

        self.color_grid_to_fen(color_grid, 'yellow', 'pink')

        print(self.board)
        print(self.board.fen())

        while(1): # game loop
            self.player_turn()
            # handle cheating
            if self.board.is_checkmate():
                break

            quit() # TODO - REMOVE

            self.ai_turn()
            if self.board.is_checkmate():
                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
        pass

    def ai_turn(self):
        # TODO
        # 1. use detect_pieces in color_analyzer
        # 2. update internal representation of board
        # 3. give to python-chess and get best move
        # 4. call move_translator with argument of best move
        pass

    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):
        for i, row in enumerate(color_grid):
            for j, (color, _, _, letter) in enumerate(row):
                if letter is not None:
                    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:
                        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("--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.show_cv, args.show_cam, args.loop, args.img_idx, args.test_img, args.save_img_as)
    game.start_game()