import cv2 import numpy as np import glob import os from sklearn import cluster from sklearn import metrics from sklearn import svm from sklearn.externals import joblib from sklearn import neural_network import heapq from datetime import datetime import utils pieces = ["rook", "knight"] colors = ['black', 'white'] def selective_search(image, use_fast=False, use_slow=False, image_name = None): # speed-up using multithreads cv2.setUseOptimized(True) cv2.setNumThreads(4) im = image img_out = im.copy() # 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) ss.switchToSingleStrategy() if (use_fast): ss.switchToSelectiveSearchFast() elif (use_slow): ss.switchToSelectiveSearchQuality() # run selective search segmentation on input image rects = ss.process() # number of region proposals to show numShowRects = 150 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 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 best_proposals.append((rect, size)) else: break height, width, channels = im.shape center_x = width // 2 center_y = (height // 2)+5 dists = [] for i in heapq.nlargest(10, best_proposals, key=lambda x: x[1]): width, height, channels = im.shape x, y, w, h = i[0] if i[1] < (width*height)*0.90 and i[1] > (width*height)*0.25: 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 dist = (center_x - box_center_x) ** 2 + (center_y - box_center_y) ** 2 dists.append([i, dist]) imCop = imOut.copy() print(image_name) best = heapq.nsmallest(1, dists, key=lambda x: x[1]) 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_right = (x + w, y + h) cropped = img_out[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]] return cropped # 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 generate_centers(number_of_clusters, sift : cv2.xfeatures2d_SIFT): features = None for piece in pieces: for color in colors: for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")): image = cv2.imread(filename) #image = selective_search(image, use_fast=True) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) kp, desc = sift.detectAndCompute(gray, None) if features is None: features = np.array(desc) else: print(f"{piece}, {color}, {filename}") features = np.vstack((features, desc)) k_means = cluster.KMeans(number_of_clusters) k_means.fit(features) return k_means.cluster_centers_ def generate_bag_of_words(image, centers, sift : cv2.xfeatures2d_SIFT): num_centers = centers.shape[0] histogram = np.zeros((1, num_centers)) gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) kp, desc = sift.detectAndCompute(gray_image, None) if not kp: return histogram distances = metrics.pairwise.pairwise_distances(desc, centers) best_centers = np.argmin(distances, axis=1) for i in best_centers: histogram[0,i] = histogram[0,i] + 1 histogram = histogram / np.sum(histogram) return histogram def do_pre_processing(): sift = cv2.xfeatures2d.SIFT_create() centers = generate_centers(8, sift) np.save("training_data/centers", centers) for piece in pieces: for color in colors: for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")): image = cv2.imread(filename) #image = selective_search(image, image_name=filename, use_fast=True) bow_features = generate_bag_of_words(image, centers, sift) np.save(f"training_data/{piece}/{color}_square/" + os.path.basename(filename), bow_features) def load_training_data(spec_piece, color): X = None Y = None for piece in pieces: piece_class = int(spec_piece == piece) for filename in glob.glob(os.path.join("training_data", piece, f"{color}_square", "*.npy")): data = np.load(filename) if X is None: X = np.array(data) Y = np.array([piece_class]) else: X = np.vstack((X, data)) Y = np.vstack((Y, [piece_class])) return X, Y def train_empty_or_piece_var(): pieces = ['empty', 'knight', 'rook'] for color in colors: X = None Y = None total_weight = 0 for piece in pieces: total_weight += len(glob.glob(os.path.join("training_images", f"{piece}", f"{color}_square", "*.png"))) current_weight = len(glob.glob(os.path.join("training_images", 'empty', f"{color}_square", "*.png"))) for piece in pieces: piece_class = int('empty' == piece) for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")): img = cv2.imread(filename) magnitude_of_var = np.linalg.norm(cv2.meanStdDev(img)[1]) if X is None: X = np.array(magnitude_of_var) Y = np.array([piece_class]) else: X = np.vstack((X, magnitude_of_var)) Y = np.vstack((Y, [piece_class])) classifier = svm.SVC(class_weight={0: current_weight, 1: total_weight - current_weight}, probability=True) classifier.fit(X, Y) joblib.dump(classifier, f"classifiers/classifier_empty_var/{color}.pkl") def train_pieces_svm(): for piece in pieces: for color in colors: # TODO: Consider removing empty from total_weights, so all classifiers do not consider empty pieces total_weights = len(glob.glob(os.path.join("training_images", "*", f"{color}_square", "*.png"))) current_weight = len(glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png"))) print(f"Trainig for piece: {piece}") X, Y = load_training_data(piece, color) classifier = svm.SVC(class_weight={0: current_weight, 1: total_weights - current_weight}, probability=True) classifier.fit(X, Y) joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl") def compute_features(training_image): sift = cv2.xfeatures2d.SIFT_create() gray_training_image = cv2.cvtColor(training_image, cv2.COLOR_BGR2GRAY) kp = sift.detect(gray_training_image) kp, desc = sift.compute(gray_training_image, kp) cv2.drawKeypoints(training_image, kp, training_image) return training_image def warp_board(camera_image, debug_image=None): #cv2.imwrite('camera_image.png', camera_image) baseline = cv2.imread("new_baseline_board.png") camera_image_gray = cv2.cvtColor(camera_image, cv2.COLOR_BGR2GRAY) baseline_gray = cv2.cvtColor(baseline, cv2.COLOR_BGR2GRAY) sift = cv2.xfeatures2d.SIFT_create() camera_image_keypoints = sift.detect(camera_image_gray, None) baseline_keypoints = sift.detect(baseline_gray, None) camera_image_keypoints, des = sift.compute(camera_image_gray, camera_image_keypoints) baseline_keypoints, des2 = sift.compute(baseline_gray, baseline_keypoints) if debug_image is not None: cv2.drawKeypoints(camera_image, keypoints=camera_image_keypoints, outImage=debug_image) cv2.imwrite('keypoints_img.jpg', camera_image) # FLANN parameters FLANN_INDEX_KDTREE = 0 index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=8) search_params = dict(checks=100) # or pass empty dictionary flann = cv2.FlannBasedMatcher(index_params,search_params) matches = flann.knnMatch(des, des2, k=2) # Need to draw only good matches, so create a mask matchesMask = [[0,0] for i in range(len(matches))] good_matches = [] # ratio test as per Lowe's paper for i,(m,n) in enumerate(matches): if m.distance < 0.55*n.distance: matchesMask[i]=[1,0] good_matches.append([m,n]) draw_params = dict(matchColor=(0,255,0), singlePointColor=(255,0,0), matchesMask=matchesMask, flags=0) img3 = cv2.drawMatchesKnn(camera_image, camera_image_keypoints, baseline, baseline_keypoints, matches, None, **draw_params) cv2.imwrite("matches.jpg", img3) # Extract location of good matches points1 = np.zeros((len(good_matches), 2), dtype=np.float32) points2 = np.zeros((len(good_matches), 2), dtype=np.float32) for i, (m, n) in enumerate(good_matches): points1[i, :] = camera_image_keypoints[m.queryIdx].pt points2[i, :] = baseline_keypoints[m.trainIdx].pt # print(len(points2)) h, mask = cv2.findHomography(points1, points2, cv2.RANSAC) height, width, channels = baseline.shape im1Reg = cv2.warpPerspective(camera_image, h, (width, height)) # cv2.imwrite('homo_pls_fuck.jpg', im1Reg) return im1Reg def get_square(warped_board, file, rank): files = "ABCDEFGH" file = files.index(file) rank = 8 - rank width, _, _ = warped_board.shape # board is square anyway side = int(width * 0.04) size = width - 2 * side square_size = size // 8 padding = 0 x1 = side + (square_size * file) x2 = x1 + square_size y1 = max(0, side + (square_size * rank) - padding) y2 = min(width, y1 + square_size + padding) square = warped_board[y1:y2, x1:x2] return square def get_squares(warped_board): result = {} for file in "ABCDEFGH": for rank in range(1, 9): square = get_square(warped_board, file, rank) result[f"{file}{rank}"] = square # cv2.imwrite(f"warped_square_{file}{rank}.png", square) return result def load_data_nn(spec_piece): X = None Y = None for piece in pieces: piece_class = int(spec_piece == piece) for filename in glob.glob(os.path.join("training_images", piece, "*", "*.png")): image = cv2.imread(filename) image = cv2.resize(image, (64, 128)) 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: 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: X = np.vstack((X, data)) Y = np.vstack((Y, [piece_class])) return (X, Y) def train_nn(): for piece in pieces: X, Y = load_data_nn(piece) classifier = neural_network.MLPClassifier(hidden_layer_sizes=64) classifier.fit(X, Y) joblib.dump(classifier, "classifiers/neural_net_" + piece + ".pkl") def letter_to_int(letter): alphabet = list('ABCDEFGH') return alphabet.index(letter) + 1 def compute_color(file, rank): if ((letter_to_int(file)+rank) % 2): return 'white' else: return 'black' def save_empty_fields(warped, skip_rank=None): alpha = "ABCDEFGH" ranks = [1, 2, 3, 4, 5, 6, 7, 8] if skip_rank is not None: ranks.remove(skip_rank) for file in alpha: for rank in ranks: square = get_square(warped, file, rank) color = compute_color(file, rank) utils.imwrite(f"training_images/empty/{color}_square/training_{file}{rank}_{datetime.utcnow().timestamp()}.png", square)