from functools import lru_cache import cv2 import runner from sklearn.externals import joblib import numpy as np import operator from matplotlib import pyplot as plt import glob import os import heapq import math from datetime import datetime 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): square_color = runner.compute_color(file, rank) y, x = np.histogram(square.ravel(), bins=256, range=[0, 256]) # TODO: Maybe img.ravel() ? for color in ['black', 'white']: empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{color}_square.pkl") prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1)) print(f"{file}{rank}, {color}: {prob[0, 1]}") if prob[0, 1] > 0.5: return 'empty' return None @lru_cache() def load_classifier(filename): print(f"loading {filename}") return joblib.load(filename) if __name__ == '__main__': board = cv2.imread("whole_boards/boards_for_empty/board_1554288951.972197_.png") warped = runner.warp_board(board) empty = 0 files = "ABCDEFGH" ranks = [1, 2, 3, 4, 5, 6, 7, 8] non_empties = [] for file in files: for rank in ranks: src = runner.get_square(warped, file, rank) #src = cv2.GaussianBlur(src, (5, 5), 0) segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.75, k=175, min_size=750) segment = segmentator.processImage(src) mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2) masked = np.ma.masked_array(src, fill_value=0) for i in range(np.max(segment)): masked.mask = mask != i y, x = np.where(segment == i) 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()}_{file}{rank}.jpg", dst) if np.max(segment) >= 2: print(f"{file}{rank} is nonempty") non_empties.append([f"{file}{rank}", src]) print(np.max(segment)) empty += 1 print(64-empty) for non_empty in non_empties: cv2.imshow(non_empty[0], non_empty[1]) cv2.waitKey(0) exit() #empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_white_square.pkl") #print(empty_classifier.predict_proba(np.array([0]*16).reshape(1, -1))[0, 1]) #exit() 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)