from functools import lru_cache 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 from sklearn.pipeline import make_pipeline from sklearn.preprocessing import StandardScaler import utils import random 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(): for square_color in ["black", "white"]: X = None Y = None for piece in ['empty', 'rook', 'knight']: for filename in glob.glob(os.path.join("training_images", f"{piece}", f"{square_color}_square", "*.png")): piece_class = 'empty' == piece img = cv2.imread(filename) y, x = np.histogram(img.ravel(), bins=256, range=[0, 256]) # TODO: Maybe img.ravel() ? if X is None: X = np.array(y) Y = np.array([piece_class]) else: X = np.vstack((X, y)) Y = np.vstack((Y, [piece_class])) classifier = make_pipeline(StandardScaler(), svm.SVC(C=10.0, gamma=0.01, probability=True)) classifier.fit(X, Y) joblib.dump(classifier, f"classifiers/classifier_empty/white_piece_on_{square_color}_square.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.045) 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 @lru_cache(maxsize=64) 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) if __name__ == '__main__': train_empty_or_piece_var()