28 lines
1.0 KiB
Python
28 lines
1.0 KiB
Python
import cv2
|
|
import numpy as np
|
|
from tensorflow.python.keras import models
|
|
|
|
from util import PIECE, Squares, Board
|
|
|
|
new_model = models.load_model('chess_model_3_pieces.h5')
|
|
#new_model.summary()
|
|
|
|
#board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
|
|
#board = cv2.imread("whole_boards/boards_for_empty/board_1554285167.655788_rank_5.png")
|
|
board = cv2.imread("whole_boards/boards_for_empty/board_1554288891.129901_rank_8.png")
|
|
|
|
|
|
def predict_board(occupied_squares: Squares) -> Board:
|
|
board = Board()
|
|
for pos, square in occupied_squares.items():
|
|
square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
|
|
width, height = square.shape
|
|
square = square / 255.0
|
|
test = new_model.predict(np.array(square).reshape((-1, width, height, 1)))
|
|
|
|
#cv2.putText(square, f"{pos} {PIECE(int(np.argmax(test)))}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255,) * 3, thickness=3)
|
|
#cv2.imwrite("lel", square)
|
|
board[pos] = PIECE(int(np.argmax(test)))
|
|
|
|
return board
|