Skip to content
Snippets Groups Projects
game.py 8.79 KiB
Newer Older
  • Learn to ignore specific revisions
  • 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
    
    zalonzo2's avatar
    zalonzo2 committed
    from board_detector import display_img
    
    import move_translator
    import cv2
    import os
    
    
    # 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
    
        def __init__(self, difficulty, color_scheme, show_cv, show_cam, loop, img_idx = 1, test_img = None, save_img_as = None):
    
    zalonzo2's avatar
    zalonzo2 committed
            self.show_cam = show_cam
    
    zalonzo2's avatar
    zalonzo2 committed
            if img_idx:
                self.img_idx = int(img_idx)
            else:
                self.img_idx = 1
    
            self.right_cut = 280
    
    zalonzo2's avatar
    zalonzo2 committed
            self.top_cut = 0
    
            self.bottom_cut = 280
    
            if not skip_camera:
                self.picam2 = Picamera2()
            else:
                self.picam2 = None
            # self.picam2 = Picamera2()
    
    zalonzo2's avatar
    zalonzo2 committed
    
    
    zalonzo2's avatar
    zalonzo2 committed
    
    
            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 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)
    
                    img_txt = self.test_img + str(self.img_idx) + '.jpg'
                    img_path = os.path.join('ocr_test_images', img_txt)
    
                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))
    
                if (self.show_cv and not self.loop):
    
                    # display_img([orig_img, img])
                    display_img([img])
    
                    answer = 0
                    while(answer != "y" and answer != "n"):
    
    zalonzo2's avatar
    zalonzo2 committed
                        answer = input("Done looping? (y/n): ")
    
            # 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(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
    
    zalonzo2's avatar
    zalonzo2 committed
        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'
    
            img_path = os.path.join('game_images', img_txt)
            metadata = self.picam2.capture_file(img_path)
    
            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 == "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")
    
    zalonzo2's avatar
    zalonzo2 committed
        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")
    
        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)