from functools import lru_cache import cv2 import runner from sklearn.externals import joblib import numpy as np import operator import glob import os import heapq import math pieces = ['rook', 'knight'] #pieces = ['rook', 'knight'] #piece_to_symbol = {'rook': 1, 'knight': 2, 'empty': 0} piece_to_symbol = {'rook': 1, 'knight': 2} colors = ['black', 'white'] def classify(image, sift : cv2.xfeatures2d_SIFT, file, rank, empty_bias=False): centers = np.load("training_data/centers.npy") probs = {'rook': {'black': 0, 'white': 0}, 'knight': {'black': 0, 'white': 0}, 'empty': {'black': 0, 'white': 0}} #probs = {'rook': 0, 'knight': 0, 'empty': 0} for piece in pieces: for color in colors: #color = runner.compute_color(file, rank) classifier = joblib.load(f"classifiers/classifier_{piece}/{color}.pkl") features = runner.generate_bag_of_words(image, centers, sift) prob = classifier.predict_proba(features) probs[piece][color] = prob[0, 1] if empty_bias: probs['empty'] *= 1.2 return probs def pred_test(file, rank, 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 = classify(mystery_image, sift, file, rank, empty_bias=empty_bias) return probs def pre_process_and_train(): runner.do_pre_processing() runner.train_pieces_svm() def build_board_from_dict(board_dict : dict): sift = cv2.xfeatures2d.SIFT_create() board = [[0]*8 for _ in range(8)] counter = 0 for idx, value in enumerate(board_dict.values()): probs = classify(value, sift) likely_piece = max(probs.items(), key=operator.itemgetter(1))[0] symbol = piece_to_symbol[likely_piece] column = idx // 8 row = (idx % 7) board[row][column] = symbol print(probs) if likely_piece != 'empty': counter += 1 print(counter) print(64/(counter-1)) return board def detect_using_nn(spec_image): probs = {'rook': 0, 'knight': 0} for piece in pieces: piece_class = piece_to_symbol[piece] win_size = (64, 64) classifier = joblib.load("classifiers/neural_net_" + piece + ".pkl") spec_image = cv2.resize(spec_image, (64, 128)) features = np.reshape(spec_image, (1, np.product(spec_image.shape))) prob = classifier.predict_proba(features) print(piece) print(prob[0,1]) def test_entire_board(): board = cv2.imread("homo_pls_fuck.jpg") warped = runner.warp_board(board) board_dict = runner.get_squares(warped) board = build_board_from_dict(board_dict) print(board) def lel_test(): # img = cv2.imread('training_images/rook/white/rook_training_D4_2.png') counter = 0 for filename in glob.glob(os.path.join("training_images", "empty", "*", "*.png")): img = cv2.imread(filename) gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) ret = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 3, 3) # binarize the image #ret, bw = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) # find connected components connectivity = 4 nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(ret, connectivity, cv2.CV_32S) sizes = stats[1:, -1] nb_components = nb_components - 1 min_size = 250 # threshhold value for objects in scene img2 = np.zeros((img.shape), np.uint8) for i in range(0, nb_components + 1): # use if sizes[i] >= min_size: to identify your objects color = np.random.randint(255, size=3) # draw the bounding rectangele around each object cv2.rectangle(img2, (stats[i][0], stats[i][1]), (stats[i][0] + stats[i][2], stats[i][1] + stats[i][3]), (0, 255, 0), 2) img2[output == i + 1] = color #print(nb_components+1) if nb_components+1 >= 4: counter += 1 print(filename) cv2.imshow("lel", img2) cv2.waitKey(0) print(counter) def selective_search(image, use_fast=False, use_slow=False): # speed-up using multithreads cv2.setUseOptimized(True) cv2.setNumThreads(4) if type(image) == str: # read image im = cv2.imread(image) else: im = image # resize image #newHeight = 200 #newWidth = int(im.shape[1] * 150 / im.shape[0]) #im = cv2.resize(im, (newWidth, newHeight)) #im = cv2.imread(image) #lel, im = cv2.threshold(im, 128, 255, cv2.THRESH_BINARY) # create Selective Search Segmentation Object using default parameters ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation() # set input image on which we will run segmentation ss.setBaseImage(im) # Switch to fast but low recall Selective Search method ss.switchToSingleStrategy() if (use_fast): ss.switchToSelectiveSearchFast() # Switch to high recall but slow Selective Search method elif (use_slow): ss.switchToSelectiveSearchQuality() # run selective search segmentation on input image rects = ss.process() #print('Total Number of Region Proposals: {}'.format(len(rects))) # number of region proposals to show numShowRects = 150 # increment to increase/decrease total number # of reason proposals to be shown increment = 1 best_proposals = [] while True: # create a copy of original image # itereate over all the region proposals for i, rect in enumerate(rects): imOut = im.copy() # draw rectangle for region proposal till numShowRects if (i < numShowRects): x, y, w, h = rect # cv2.rectangle(imOut, (x, y), (x + w, y + h), (0, 255, 0), 1, cv2.LINE_AA) # size = (max(w, x) - min(w, x)) * ((max(h, y) - min(h, y))) top_left = (x,y) bottom_left = (x, y+h) top_right = (x+w, y) bottom_right = (x+w, y+h) rect_width = bottom_right[0] - bottom_left[0] rect_height = bottom_right[1] - top_right[1] size = rect_width * rect_height #print(f"({x}, {y}), ({w}, {h})\n Of size: { size }") #cv2.rectangle(imOut, (x, y), (x + w, y + h), (0, 255, 0), 1, cv2.LINE_AA) #cv2.imshow("lel", imOut) #cv2.waitKey(0) best_proposals.append((rect, size)) #if size > biggest_size: # biggest_rect = (x, y, w, h) # biggest_size = size # print(f"New biggest: \n({x}, {y}), ({w}, {h})\nOf size: {biggest_size}") else: break height, width, channels = im.shape center_x = width // 2 center_y = (height // 2)+5 dists = [] #print(f"Amount of best proposals:\n{len(best_proposals)}") #print(f"lel: {len(heapq.nlargest(10, best_proposals, key=lambda x: x[1]))}") for i in heapq.nlargest(10, best_proposals, key=lambda x: x[1]): width, height, channels = im.shape #print(width * height) #print(i[1]) x, y, w, h = i[0] if i[1] <= (width*height)*0.8 and i[1] > (width*height)*0.25: imCop = imOut.copy() #cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 2, cv2.LINE_AA) #cv2.imshow("lel", imCop) #cv2.waitKey(0) #cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 4, cv2.LINE_AA) top_left = (x,y) bottom_left = (x, y+h) top_right = (x+w, y) bottom_right = (x+w, y+h) box_center_x = (top_left[0]+bottom_left[0]+top_right[0]+bottom_right[0]) // 4 box_center_y = (top_left[1]+bottom_left[1]+top_right[1]+bottom_right[1]) // 4 #print(f"{box_center_x}, {box_center_y}, {center_x}, {center_y}") dist = (center_x - box_center_x) ** 2 + (center_y - box_center_y) ** 2 print(dist) dists.append([i, dist]) cv2.drawMarker(imCop, position=(x+w, h+y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x+w, y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x, y+h), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3) cv2.drawMarker(imCop, position=(center_x, center_y), color=(0, 0, 255), thickness=3) #cv2.imshow("lel", imCop) #cv2.waitKey(0) #print("-------"*5) for pls in dists: imCop = imOut.copy() x, y, w, h = pls[0][0] #print(x,y,w,h) #print(pls[1]) top_left = (x, y) bottom_left = (x, y + h) top_right = (x + w, y) bottom_right = (x + w, y + h) cv2.drawMarker(imCop, position=(x + w, h + y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x + w, y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x, y + h), color=(255, 0, 0), thickness=3) box_center_x = (top_left[0] + bottom_left[0] + top_right[0] + bottom_right[0]) // 4 box_center_y = (top_left[1] + bottom_left[1] + top_right[1] + bottom_right[1]) // 4 cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3) cv2.drawMarker(imCop, position=(center_y, center_x), color=(0, 0, 255), thickness=3) cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 2, cv2.LINE_AA) #cv2.imshow("lel", imCop) #cv2.waitKey(0) imCop = imOut.copy() best = heapq.nsmallest(1, dists, key=lambda x: x[1]) if (len(best) == 0): return ((0, 0), (0, height), (width, 0), (width, height)) x, y, w, h = best[0][0][0] cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 4, cv2.LINE_AA) top_left = (x, y) bottom_left = (x, y + h) top_right = (x + w, y) bottom_right = (x + w, y + h) cv2.drawMarker(imCop, position=(x + w, h + y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x + w, y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3) cv2.drawMarker(imCop, position=(x, y + h), color=(255, 0, 0), thickness=3) box_center_x = (top_left[0] + bottom_left[0] + top_right[0] + bottom_right[0]) // 4 box_center_y = (top_left[1] + bottom_left[1] + top_right[1] + bottom_right[1]) // 4 cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3) cv2.drawMarker(imCop, position=(center_x, center_y), color=(0, 0, 255), thickness=3) #cv2.imshow("lel", imCop) #cv2.waitKey(0) return (top_left, bottom_left, top_right, bottom_right) # show output cv2.imshow("Output", imOut) # record key press k = cv2.waitKey(0) & 0xFF # m is pressed if k == 109: # increase total number of rectangles to show by increment numShowRects += increment # l is pressed elif k == 108 and numShowRects > increment: # decrease total number of rectangles to show by increment numShowRects -= increment # q is pressed elif k == 113: break # close image show window cv2.destroyAllWindows() def predict(square, file, rank): color = runner.compute_color(file, rank) empty_var_classifier = load_classifier(f"classifiers/classifier_empty_var/{color}.pkl") magnitude_of_var = np.linalg.norm(cv2.meanStdDev(square)[1]) prob = empty_var_classifier.predict_proba(np.array(magnitude_of_var).reshape(-1, 1)) print(prob[0, 1]) if (prob[0, 1]) > 0.5: return 'empty' return None @lru_cache() def load_classifier(filename): return joblib.load(filename) if __name__ == '__main__': board = cv2.imread("whole_boards/board_102_1554110461.608167_.png") warped = runner.warp_board(board) files = "ABCDEFGH" ranks = [1,2,3,4,5,6,7,8] counter = 0 for file in files: for rank in ranks: square = runner.get_square(warped, file, rank) if predict(square, file, rank) == 'empty': counter += 1 print(counter) exit() square = runner.get_square(warped, "D", 2) gray_square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY) print(cv2.meanStdDev(gray_square)[1]) print(cv2.meanStdDev(square)[1]) cv2.imshow("square", square) cv2.waitKey(0) print(pred_test("C", 2, square)) sift: cv2.xfeatures2d_SIFT = cv2.xfeatures2d.SIFT_create() gray = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY) kp, desc = sift.detectAndCompute(gray, None) cv2.drawKeypoints(square, kp, square) cv2.imshow("kp", square) cv2.waitKey(0) exit() board = cv2.imread("whole_boards/board_202_1554154094.001122_.png") runner.fetch_empty_fields(board) exit() warped = runner.warp_board(board) counter = 0 #square = runner.get_square(warped, "A", 3) #top_left, bottom_left, top_right, bottom_right = selective_search(square, use_fast=True) #cropped = square[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]] for file in files: for rank in ranks: square = runner.get_square(warped, file, rank) top_left, bottom_left, top_right, bottom_right = selective_search(square, use_fast=True) cropped = square[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]] rect_width = bottom_right[0] - bottom_left[0] rect_height = bottom_right[1] - top_right[1] size = rect_width * rect_height square_height, square_width, channels = square.shape empty_bias = (size == square_height*square_width) if size == square_height*square_width: print(f"{file}{rank} is likely empty") res = pred_test(file, rank, mystery_image=square, empty_bias=empty_bias) print(res) if (max(res.items(), key=operator.itemgetter(1))[0] == 'empty'): counter += 1 print(f"Amount of empty fields: {counter}") #print("Non-cropped:\t",pred_test(square)) #print("Cropped:\t",pred_test(cropped)) #cv2.imshow("square", square) #cv2.waitKey(0) #runner.do_pre_processing() #runner.train() #img = "warped_square_B5.png" #detect_using_nn(img) #selective_search("training_images/empty/colorless/warped_square_A6.png", use_fast=True) #selective_search("warped_square_B5.png", use_fast=True) img = "training_images/rook/white/rook_training_D4_7.png" #img = "training_images/rook/white_square/rook_training_E4_10.png" #img = "training_images/knight/white_square/training_D5_134.png" #top_left, bottom_left, top_right, bottom_right = selective_search(img, use_fast=True) #cropped = cv2.imread(img)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]] #cv2.imshow("output", cropped) #print(pred_test(cropped)) #cv2.waitKey(0) #lel_test() # test_entire_board() #board = [[0, 0, 1, 2, 0, 0, 0, 2], [0, 1, 2, 2, 1, 0, 0, 1], [0, 0, 0, 0, 1, 0, 2, 0], [0, 2, 2, 1, 1, 2, 2, 0], [0, 1, 0, 0, 1, 2, 0, 0], [0, 0, 0, 0, 0, 2, 2, 0], [0, 0, 0, 2, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0]] #for i in board: # print(i) #warped = cv2.imread("homo_pls_fuck.jpg") #square = runner.get_square(warped, "D", 4) #print(pred_test(square)) #cv2.imshow("lel", square) #cv2.waitKey(0)