Skip to content
Snippets Groups Projects
game.py 3.48 KiB
Newer Older
import argparse
import chess
import chess.pgn
from board_detector import find_board_and_pieces
from board_detector import init_show_cv
import color_analyzer
import move_translator
import cv2
import os

class ChessGame:
    def __init__(self, difficulty, show_cv, test_img = None):
        self.board = chess.Board()
        self.difficulty = difficulty
        self.show_cv = show_cv
        self.test_img = test_img

    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_show_cv(self.show_cv)

        if (self.test_img):
            img_path = os.path.join('test_images', self.test_img)
            orig_img = cv2.imread(img_path)
            orig_img = cv2.imread('test_images/board1.jpg') # TODO - CHANGE TO MAKE IT RECEIVE INPUT FROM CAMERA
        h,w,c = orig_img.shape
        print(h, w)
        cropped_img = orig_img
        cropped_img = orig_img[0:h, int(w/2 - h/2 - 60):int(w/2 + h/2 + 100)]
        img = cv2.resize(cropped_img, (512, 512))
        # img = cv2.resize(cropped_img, (int(w/5), int(h/5)))
        # img = orig_img

        if (self.show_cv):
            cv2.imshow('Original Image Before Processing', img)
            cv2.waitKey(0)
            cv2.destroyAllWindows()


        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

        # while(1): # loop until legal move has been made
        #     cur_state = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
        #     potential_next_state = "rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq - 0 1"
        #     board = chess.Board(cur_state)
        #     next_board = chess.Board(potential_next_state)
        #     move = next_board.peek() # NOTHING ON STACK SO DOESN'T WORK
        #     if next_board.is_legal(move):
        #         board.set_fen(potential_next_state)
        #         print("Move executed successfully.")
        #     else:
        #         print("Illegal move. Move not executed.")

        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

    
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("--test_img", help="If specified, will use said image in test_images folder rather than camera input")
    args = parser.parse_args()

    game = ChessGame(args.difficulty, args.show_cv, args.test_img)
    game.start_game()