diff --git a/README.md b/README.md index 18b383a1..dc1c4c5d 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ ## Test Adapter: ```bash -base64 --wrap=0 < whole_boards/board_118_1554110522.620303_.png| python3 adapter.py +base64 --wrap=0 < whole_boards/boards_for_empty/board_1554285167.655788_rank_5.png | python3 adapter.py ``` ## Build OpenCV diff --git a/adapter.py b/adapter.py index 199471a6..f2282ea1 100644 --- a/adapter.py +++ b/adapter.py @@ -5,21 +5,29 @@ import sys import cv2 import numpy as np -from runner import find_homography - +from main import find_occupied_squares +from runner import find_homography, warp_board # Load base64 encoded image from stdin +from tensor_classifier import predict_board + stdin = sys.stdin.readline() stdin_decoded = base64.b64decode(stdin) img_array = np.frombuffer(stdin_decoded, dtype=np.uint8) camera_img = cv2.imdecode(img_array, flags=cv2.COLOR_BGR2RGB) camera_img = cv2.cvtColor(camera_img, cv2.COLOR_BGR2RGB) -# Find keypoints in image and pass them back to unity + +# def do_everything: homography = find_homography(camera_img) +warped_board = warp_board(camera_img, homography) +occupied_squares = find_occupied_squares(warped_board) +board = predict_board(occupied_squares) + # Finally, output to stdout for unity to read result = { "homography": homography.tolist(), + "board": board.to_array, } print(json.dumps(result)) diff --git a/classifiers/classifier_empty/white_piece_on_black_square.pkl b/classifiers/classifier_empty/white_piece_on_black_square.pkl index 7cbdd741..ca569b8c 100644 Binary files a/classifiers/classifier_empty/white_piece_on_black_square.pkl and b/classifiers/classifier_empty/white_piece_on_black_square.pkl differ diff --git a/classifiers/classifier_empty/white_piece_on_white_square.pkl b/classifiers/classifier_empty/white_piece_on_white_square.pkl index 18436524..3807cfff 100644 Binary files a/classifiers/classifier_empty/white_piece_on_white_square.pkl and b/classifiers/classifier_empty/white_piece_on_white_square.pkl differ diff --git a/main.py b/main.py index f5f80e3c..6db9d81e 100644 --- a/main.py +++ b/main.py @@ -102,6 +102,7 @@ def predict_empty(square: np.ndarray, position: POSITION) -> PIECE: return None + def remove_most_empties(warped): empty = 0 non_empties = [] @@ -123,12 +124,6 @@ def remove_most_empties(warped): y, x = np.where(segment == i) - - if position == POSITION.E8: - print(np.max(segment)) - print(len(y)) - - pls.append(len(y)) top, bottom, left, right = min(y), max(y), min(x), max(x) @@ -146,13 +141,13 @@ def remove_most_empties(warped): return non_empties -def find_occupied_squares(warped: np.ndarray) -> List[Tuple[POSITION, np.ndarray]]: +def find_occupied_squares(warped: np.ndarray) -> Squares: non_empties = remove_most_empties(warped) - completely_non_empties = [] + completely_non_empties = {} for position, square in non_empties: if predict_empty(square, position) != PIECE.EMPTY: - completely_non_empties.append((position, square)) + completely_non_empties[position] = square return completely_non_empties diff --git a/opencv_video.py b/opencv_video.py index 64bef30b..94eaf711 100644 --- a/opencv_video.py +++ b/opencv_video.py @@ -10,11 +10,11 @@ cap = cv2.VideoCapture(0) color = COLOR.BLACK rank = RANK.EIGHT pieces = { - PIECE.rook: [POSITION(FILE.A, rank), POSITION(FILE.F, rank)], - PIECE.knight: [POSITION(FILE.E, rank), POSITION(FILE.H, rank)], - PIECE.bishop: [POSITION(FILE.C, rank), POSITION(FILE.D, rank)], - PIECE.queen: [POSITION(FILE.B, rank)], - PIECE.king: [POSITION(FILE.G, rank)], + PIECE.rook: [POSITION((FILE.A, rank)), POSITION((FILE.F, rank))], + PIECE.knight: [POSITION((FILE.E, rank)), POSITION((FILE.H, rank))], + PIECE.bishop: [POSITION((FILE.C, rank)), POSITION((FILE.D, rank))], + PIECE.queen: [POSITION((FILE.B, rank))], + PIECE.king: [POSITION((FILE.G, rank))], } while True: diff --git a/runner.py b/runner.py index a71c7f28..9ff20a50 100644 --- a/runner.py +++ b/runner.py @@ -89,6 +89,7 @@ def train_empty_or_piece_hist() -> None: X = [] Y = [] for piece in OUR_PIECES + (PIECE.EMPTY,): + print(f"training training_images/{piece}/{square_color}_square/*.png") for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"): img = cv2.imread(filename) y, x = np.histogram(img.ravel(), bins=32, range=[0, 256]) @@ -209,7 +210,8 @@ def find_homography(camera_image: np.ndarray, def warp_board(camera_image: np.ndarray, homography: np.ndarray = None, debug=False) -> np.ndarray: baseline = cv2.imread(str(here.joinpath("new_baseline_board.png"))) - homography = homography or find_homography(camera_image, baseline, debug=debug) + if homography is None: + homography = find_homography(camera_image, baseline, debug=debug) height, width, channels = baseline.shape return cv2.warpPerspective(camera_image, homography, (width, height)) @@ -289,4 +291,5 @@ def train_nn(): if __name__ == '__main__': - train_nn() + #train_nn() + train_empty_or_piece_hist() diff --git a/tensor_classifier.py b/tensor_classifier.py index f86fa945..dad7c7ad 100644 --- a/tensor_classifier.py +++ b/tensor_classifier.py @@ -1,71 +1,27 @@ -#from __future__ import absolute_import, division, print_function, unicode_literals - -import tensorflow as tf -from tensorflow.python.keras import datasets, layers, models -import glob -import numpy as np import cv2 +import numpy as np +from tensorflow.python.keras import models -import runner -from main import find_occupied_squares -from util import POSITION +from util import PIECE, Squares, Board new_model = models.load_model('chess_model_3_pieces.h5') -new_model.summary() +#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") -warped = runner.warp_board(board) - -occupied = find_occupied_squares(warped) - -pos_1 = POSITION.C5 -pos_2 = POSITION.D5 -pos_3 = POSITION.G5 -pos_4 = POSITION.H5 - -square_1 = runner.get_square(warped, pos_1) -square_2 = runner.get_square(warped, pos_2) -square_3 = runner.get_square(warped, pos_3) -square_4 = runner.get_square(warped, pos_4) - -square_1 = cv2.cvtColor(square_1, cv2.COLOR_BGR2GRAY) -square_2 = cv2.cvtColor(square_2, cv2.COLOR_BGR2GRAY) -square_3 = cv2.cvtColor(square_3, cv2.COLOR_BGR2GRAY) -square_4 = cv2.cvtColor(square_4, cv2.COLOR_BGR2GRAY) - -width, height = square_1.shape -square_1 = square_1 / 255.0 -square_2 = square_2 / 255.0 -square_3 = square_3 / 255.0 -square_4 = square_4 / 255.0 - - -pieces = ['knight', 'rook', 'bishop'] - -for pos, square in occupied: - 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))) - text_color = 255 - cv2.putText(square, f"{pos} {pieces[int(np.argmax(test))]}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, - color=(text_color,) * 3, thickness=3) - cv2.imwrite("lel", square) - -cv2.waitKey(0) - -""" -for pos, square in [(pos_1, square_1), (pos_2, square_2), (pos_3, square_3), (pos_4, square_4)]: - test = new_model.predict(np.array(square).reshape((-1, width, height, 1))) - print(f"{pos}: {np.argmax(test)}") - text_color = 255 - cv2.putText(square, f"{pos} {pieces[int(np.argmax(test))]}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, - color=(text_color,) * 3, thickness=3) - cv2.imshow(f"{pos}", square) - """ +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 diff --git a/tmp/tensor.py b/tmp/tensor.py index 3ef49a3d..7d7afdc9 100644 --- a/tmp/tensor.py +++ b/tmp/tensor.py @@ -1,52 +1,35 @@ from __future__ import absolute_import, division, print_function, unicode_literals -import tensorflow as tf import glob -import numpy as np + import cv2 +import numpy as np +import tensorflow as tf -#(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data() +# exit() +from util import OUR_PIECES -#print(train_images[0]) +# (train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data() +# print(train_images[0]) -#exit() training_img = [] training_labels = [] test_img = [] test_labels_ = [] +for piece in OUR_PIECES: + # training set + for _ in range(10): + for filename in glob.glob(f"../training_images/{piece}/*_square/*.png")[:-50]: + training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) + training_labels.append(piece) -for _ in range(10): - for filename in glob.glob(f"../training_images/rook/*_square/*.png")[:-50]: - training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) - training_labels.append(1) - - -for _ in range(10): - for filename in glob.glob(f"../training_images/knight/*_square/*.png")[:-50]: - training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) - training_labels.append(0) - -for _ in range(10): - for filename in glob.glob(f"../training_images/bishop/*_square/*.png")[:-50]: - training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) - training_labels.append(2) - -for _ in range(5): - for filename in glob.glob(f"../training_images/rook/*_square/*.png")[-50:]: - test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) - test_labels_.append(1) - -for _ in range(5): - for filename in glob.glob(f"../training_images/bishop/*_square/*.png")[-50:]: - test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) - test_labels_.append(2) - -for _ in range(5): - for filename in glob.glob(f"../training_images/knight/*_square/*.png")[-50:]: - test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) - test_labels_.append(0) + # test set + for _ in range(5): + for filename in glob.glob(f"../training_images/{piece}/*_square/*.png")[-50:]: + test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY)) + test_labels_.append(piece) width, height = training_img[0].shape diff --git a/util.py b/util.py index 263a2447..3d6e442c 100644 --- a/util.py +++ b/util.py @@ -3,7 +3,7 @@ from __future__ import annotations from enum import Enum from functools import lru_cache from pathlib import Path -from typing import NamedTuple, Dict, Tuple +from typing import NamedTuple, Dict, Tuple, List import cv2 import numpy as np @@ -19,23 +19,24 @@ class COLOR(Enum): class PIECE(Enum): - PAWN = "pawn" - ROOK = "rook" - KNIGHT = "knight" - BISHOP = "bishop" - QUEEN = "queen" - KING = "king" - EMPTY = "empty" + KNIGHT = 0 + ROOK = 1 + BISHOP = 2 + PAWN = 3 + QUEEN = 4 + KING = 5 + EMPTY = 6 def __str__(self) -> str: - return self.value + return self.name.lower() PieceAndColor = Tuple[PIECE, COLOR] OUR_PIECES = ( - PIECE.ROOK, PIECE.KNIGHT, + PIECE.ROOK, + PIECE.BISHOP, ) @@ -85,6 +86,11 @@ Squares = Dict[POSITION, np.ndarray] class Board(Dict[POSITION, PIECE]): """Board is a dict mapping positions to a piece, i.e. a board configuration after all image processing""" + @property + def to_array(self) -> List[List[int]]: + return [[self.get(POSITION((file, rank)), PIECE.EMPTY).value for file in FILE] + for rank in RANK] + def imwrite(*args, **kwargs): Path(args[0]).parent.mkdir(parents=True, exist_ok=True)