import sys import warnings import cv2 import matplotlib.pyplot as plt import numpy as np from sklearn.exceptions import DataConversionWarning import runner from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor, OUR_PIECES warnings.filterwarnings(action='ignore', category=DataConversionWarning) np.set_printoptions(threshold=sys.maxsize) def identify_piece(image: np.ndarray, position: POSITION, sift: cv2.xfeatures2d_SIFT) -> PieceAndColor: centers = np.load("training_data/centers.npy") best = 0 probs = {p.name: {} for p in OUR_PIECES} best_piece = best_color = None 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 if prob[0, 1] > best: best = prob[0, 1] best_piece, best_color = piece, color #print(probs) return best_piece, best_color def pred_test(position: POSITION, mystery_image=None, empty_bias=False): sift = cv2.xfeatures2d.SIFT_create() if mystery_image is None: mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png") probs = identify_piece(mystery_image, position, sift) return probs def pre_process_and_train() -> None: runner.do_pre_processing() runner.train_pieces_svm() def build_board_from_squares(squares: Squares) -> Board: sift = cv2.xfeatures2d.SIFT_create() board = Board() counter = 0 for position, square in squares.values(): likely_piece = identify_piece(square, position, sift) board[position] = likely_piece if likely_piece != PIECE.EMPTY: counter += 1 print(counter) print(64/(counter-1)) return board def test_entire_board() -> None: board_img = cv2.imread("homo_pls_fuck.jpg") warped = runner.warp_board(board_img) squares = runner.get_squares(warped) board = build_board_from_squares(squares) print(board) def predict_empty(square: np.ndarray, position: POSITION) -> bool: y, x = np.histogram(square.ravel(), bins=32, range=[0, 256]) left, right = x[:-1], x[1:] X = np.array([left, right]).T.flatten() Y = np.array([y, y]).T.flatten() area = sum(np.diff(x) * y) plt.plot(X, Y) plt.xlabel(f"{position}") #plt.show() empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{position.color}_square.pkl") prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1)) #print(f"{position}, {position.color}: {prob[0, 1]}") return prob[0, 1] > 0.95 def remove_most_empties(warped): empty = 0 non_empties = [] for position in POSITION: counter = 0 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) mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2) masked = np.ma.masked_array(src, fill_value=0) pls = [] for i in range(np.max(segment)): masked.mask = mask != i y, x = np.where(segment == i) 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"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.9469 for x in pls])): #print(f"{position} is nonempty") non_empties.append([position, img_src]) empty += 1 #print(64 - empty) return non_empties def find_occupied_squares(warped: np.ndarray) -> Squares: non_empties = remove_most_empties(warped) completely_non_empties = {} for position, square in non_empties: if not predict_empty(square, position): completely_non_empties[position] = square return completely_non_empties if __name__ == '__main__': #runner.train_pieces_svm() 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) knight_out = cv2.Canny(knight_square, 50, 55, L2gradient=True) 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() for position, square in occupied: print("---"*15) piece, color = identify_piece(square, position, sift) print(f"{piece} on {position}") text_color = 255 if color == COLOR.WHITE else 0 cv2.putText(square, f"{position} {piece.name}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(text_color,)*3, thickness=3) cv2.imshow(f"{position}", square) cv2.waitKey(0)