diff --git a/classifiers/classifier_empty/.DS_Store b/classifiers/classifier_empty/.DS_Store new file mode 100644 index 00000000..c098dba6 Binary files /dev/null and b/classifiers/classifier_empty/.DS_Store differ diff --git a/classifiers/classifier_empty/white_piece_on_black_square.pkl b/classifiers/classifier_empty/white_piece_on_black_square.pkl new file mode 100644 index 00000000..7cbdd741 Binary files /dev/null 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 new file mode 100644 index 00000000..18436524 Binary files /dev/null and b/classifiers/classifier_empty/white_piece_on_white_square.pkl differ diff --git a/classifiers/classifier_empty_var/black.pkl b/classifiers/classifier_empty_var/black.pkl new file mode 100644 index 00000000..295dbfbf Binary files /dev/null and b/classifiers/classifier_empty_var/black.pkl differ diff --git a/classifiers/classifier_empty_var/white.pkl b/classifiers/classifier_empty_var/white.pkl new file mode 100644 index 00000000..7d80bc8c Binary files /dev/null and b/classifiers/classifier_empty_var/white.pkl differ diff --git a/classifiers/classifier_knight.pkl b/classifiers/classifier_knight.pkl new file mode 100644 index 00000000..ac81e714 Binary files /dev/null and b/classifiers/classifier_knight.pkl differ diff --git a/classifiers/classifier_knight/black.pkl b/classifiers/classifier_knight/black.pkl new file mode 100644 index 00000000..f250e8e7 Binary files /dev/null and b/classifiers/classifier_knight/black.pkl differ diff --git a/classifiers/classifier_knight/white.pkl b/classifiers/classifier_knight/white.pkl new file mode 100644 index 00000000..afb83494 Binary files /dev/null and b/classifiers/classifier_knight/white.pkl differ diff --git a/classifiers/classifier_rook.pkl b/classifiers/classifier_rook.pkl new file mode 100644 index 00000000..ace5777e Binary files /dev/null and b/classifiers/classifier_rook.pkl differ diff --git a/classifiers/classifier_rook/black.pkl b/classifiers/classifier_rook/black.pkl new file mode 100644 index 00000000..6769df8d Binary files /dev/null and b/classifiers/classifier_rook/black.pkl differ diff --git a/classifiers/classifier_rook/white.pkl b/classifiers/classifier_rook/white.pkl new file mode 100644 index 00000000..646ef1f7 Binary files /dev/null and b/classifiers/classifier_rook/white.pkl differ diff --git a/main.py b/main.py index 49ff0b31..f5f80e3c 100644 --- a/main.py +++ b/main.py @@ -24,9 +24,15 @@ def identify_piece(image: np.ndarray, position: POSITION, sift: cv2.xfeatures2d_ for piece in OUR_PIECES: for color in COLOR: #color = runner.compute_color(file, rank) + #classifier = load_classifier(f"classifiers/neural_net_{piece}/{color}.pkl") classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl") features = runner.generate_bag_of_words(image, centers, sift) prob = classifier.predict_proba(features) + + #image = cv2.resize(image, (172, 172)) + #data = np.reshape(image, (1, np.product(image.shape))) + #prob = classifier.predict_proba(data) + probs[piece.name][color.name] = prob[0, 1] print(f"{piece}, {color}, {prob[0, 1]}") #if prob[0, 1] > best and color == position.color: # can only be best if correct color. Iterating through both colors for debugging only @@ -102,9 +108,9 @@ def remove_most_empties(warped): for position in POSITION: counter = 0 - src = runner.get_square(warped, position) - width, height, _ = src.shape - src = src[width // 25:, height // 25:] + img_src = runner.get_square(warped, position) + width, height, _ = img_src.shape + src = img_src[width // 25:, height // 25:] # src = src[:-width//200, :-height//200] segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.8, k=150, min_size=700) segment = segmentator.processImage(src) @@ -116,16 +122,23 @@ def remove_most_empties(warped): masked.mask = mask != i 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) dst = masked.filled()[top: bottom + 1, left: right + 1] - cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{position}.png", dst) + cv2.imwrite(f"tmp_seg/segment_{datetime.utcnow().timestamp()}_{position}.png", dst) if np.max(segment) > 0 and not np.all([x < (164 ** 2) * 0.2 for x in pls]) and ( - np.max(segment) >= 3 or np.all([x < (164 ** 2) * 0.942 for x in pls])): + np.max(segment) >= 3 or np.all([x < (164 ** 2) * 0.9469 for x in pls])): #print(f"{position} is nonempty") - non_empties.append([position, src]) + non_empties.append([position, img_src]) empty += 1 #print(64 - empty) @@ -148,21 +161,29 @@ if __name__ == '__main__': #runner.train_pieces_svm() - #board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png") - board = cv2.imread("whole_boards/boards_for_empty/board_1554286515.323962_rank_3.png") + board = cv2.imread("whole_boards/boards_for_empty/board_1554288891.129901_rank_8.png") + #board = cv2.imread("whole_boards/boards_for_empty/board_1554286515.323962_rank_3.png") warped = runner.warp_board(board) + tmp = find_occupied_squares(warped) + for pos, square in tmp: + cv2.imshow(f"{pos}", square) + cv2.waitKey(0) + exit() + + """ rook_square = runner.get_square(warped, POSITION.H3) knight_square = runner.get_square(warped, POSITION.D3) cv2.imshow("lel", rook_square) cv2.imshow("lil", knight_square) - rook_out = cv2.Canny(rook_square, 50, 55, L2gradient=True) + #rook_out = cv2.Canny(rook_square, 50, 55, L2gradient=True) knight_out = cv2.Canny(knight_square, 50, 55, L2gradient=True) - cv2.imshow("lal", rook_out) - cv2.imshow("lul", knight_out) + knight_out_l = cv2.Canny(knight_square, 50, 55, L2gradient=False) + cv2.imshow("lal", knight_out) + cv2.imshow("lul", knight_out_l) cv2.waitKey(0) exit() - + """ occupied = find_occupied_squares(warped) sift = cv2.xfeatures2d.SIFT_create() diff --git a/runner.py b/runner.py index 94e3536d..a69ee945 100644 --- a/runner.py +++ b/runner.py @@ -9,6 +9,7 @@ from sklearn import cluster, metrics, svm from sklearn.externals import joblib from sklearn.pipeline import make_pipeline from sklearn.preprocessing import StandardScaler +from sklearn import neural_network from util import RANK, POSITION, imwrite, PIECE, COLOR, Squares, OUR_PIECES @@ -108,11 +109,28 @@ def train_pieces_svm() -> None: current_weight = len(glob.glob(f"training_images/{piece}/{color}_square/*.png")) print(f"Training for piece: {piece}") X, Y = load_training_data(piece, color) - classifier = svm.SVC(C=10, gamma=0.01, class_weight={0: 45, 1: 1}, probability=True) + classifier = svm.SVC(C=10, gamma=0.01, class_weight={0: 15, 1: 0.8}, probability=True) classifier.fit(X, Y) joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl") +def train_pieces_svm_canny() -> None: + for square_color in COLOR: + X = [] + Y = [] + for piece in OUR_PIECES + (PIECE.EMPTY,): + 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]) + X.append(y) + Y.append(piece == PIECE.EMPTY) + + classifier = make_pipeline(StandardScaler(), + svm.SVC(C=10.0, gamma=0.01, probability=True)) + classifier.fit(np.array(X), np.array(Y)) + joblib.dump(classifier, f"classifiers/classifier_empty/white_piece_on_{square_color}_square.pkl") + + def warp_board(camera_image, debug_image=None) -> np.ndarray: baseline = cv2.imread("new_baseline_board.png") @@ -181,10 +199,10 @@ def get_square(warped_board: np.ndarray, position: POSITION) -> np.ndarray: side = int(width * 0.045) size = width - 2 * side square_size = size // 8 - padding = 0 + padding = 2 x1 = side + (square_size * (position.file - 1)) - x2 = x1 + square_size + x2 = x1 + square_size + padding y1 = max(0, side + (square_size * (8 - position.rank)) - padding) # 8 - rank because chessboard is from 8 to 1 y2 = min(width, y1 + square_size + padding) @@ -207,5 +225,47 @@ def save_empty_fields(warped_board: np.ndarray, skip_rank: RANK = None) -> None: imwrite(f"training_images/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square) +def load_data_nn(spec_piece, color): + X = None + Y = None + for piece in OUR_PIECES: + piece_class = int(spec_piece == piece) + + for filename in glob.glob(f"training_images/{piece}/{color}_square/*.png"): + image = cv2.imread(filename) + data = np.reshape(image, (1, np.product(image.shape))) + if X is None: + if piece_class == 1: + for _ in range(10): + X = np.array(data) + Y = np.array([piece_class]) + else: + for _ in range(5): + X = np.array(data) + Y = np.array([piece_class]) + else: + if piece_class == 1: + for _ in range(10): + X = np.vstack((X, data)) + Y = np.vstack((Y, [piece_class])) + else: + for _ in range(5): + X = np.vstack((X, data)) + Y = np.vstack((Y, [piece_class])) + return X, Y + + +def train_nn(): + for piece in OUR_PIECES: + for color in COLOR: + X, Y = load_data_nn(piece, color) + + classifier = neural_network.MLPClassifier(hidden_layer_sizes=256) + + classifier.fit(X, Y) + + joblib.dump(classifier, f"classifiers/neural_net_{piece}/{color}.pkl") + + if __name__ == '__main__': - train_empty_or_piece_hist() + train_nn() diff --git a/tensor_classifier.py b/tensor_classifier.py new file mode 100644 index 00000000..8a28ea10 --- /dev/null +++ b/tensor_classifier.py @@ -0,0 +1,72 @@ +from __future__ import absolute_import, division, print_function, unicode_literals + +import tensorflow as tf +import tensorflow_datasets as tfds +from tensorflow.python.keras import datasets, layers, models +import glob +import numpy as np +import cv2 + +import runner +from main import find_occupied_squares +from util import POSITION + +new_model = models.load_model('test_chess_model.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") + + +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'] + +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.imshow(f"{pos}", 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) + """ + + diff --git a/test_chess_model.h5 b/test_chess_model.h5 new file mode 100644 index 00000000..7b300df5 Binary files /dev/null and b/test_chess_model.h5 differ