import glob import os import time from datetime import datetime from pathlib import Path from typing import Tuple import copyreg import cv2 import numpy as np from sklearn import cluster, metrics, svm, neural_network from sklearn.externals import joblib from sklearn.pipeline import make_pipeline from sklearn.preprocessing import StandardScaler from util import RANK, POSITION, imwrite, PIECE, COLOR, Squares, OUR_PIECES here: Path = Path(__file__).parent BASELINE = cv2.imread(str(here.joinpath("new_baseline_board.png"))) BASELINE_GRAY = cv2.cvtColor(BASELINE, cv2.COLOR_BGR2GRAY) SIFT = cv2.xfeatures2d.SIFT_create() BASELINE_KEYPOINTS = SIFT.detect(BASELINE_GRAY) BASELINE_KEYPOINTS, BASELINE_DES = SIFT.compute(BASELINE_GRAY, BASELINE_KEYPOINTS) def generate_centers(number_of_clusters, sift: cv2.xfeatures2d_SIFT): features = None for piece in OUR_PIECES: for color in COLOR: for filename in glob.glob(f"training_images/{piece}/{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) print(f"{piece}, {color}, {filename}") if features is None: features = np.array(desc) else: print(f"{piece}, {color}, {filename}") features = np.vstack((features, desc)) features = np.array(features) 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: # TODO: Could do this way faster in one line with numpy somehow histogram[0, i] += + 1 return histogram / np.sum(histogram) def do_pre_processing() -> None: sift = cv2.xfeatures2d.SIFT_create() centers = generate_centers(8, sift) np.save("training_data/centers", centers) for piece in OUR_PIECES: for color in COLOR: for filename in glob.glob(f"training_images/{piece}/{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(piece: PIECE, color: COLOR) -> Tuple[np.array, np.array]: X = [] Y = [] for p in OUR_PIECES: for filename in glob.glob(f"training_data/{piece}/{color}_square/*.npy"): data = np.load(filename) X.append(data[0]) Y.append(p == piece) return np.array(X), np.array(Y) def train_empty_or_piece_hist() -> None: for square_color in COLOR: X = [] Y = [] for piece in OUR_PIECES + (PIECE.EMPTY,): print(f"training training_images/{piece}/{square_color}_square/*.png") for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"): img = cv2.imread(filename) y, x = np.histogram(img.ravel(), bins=32, range=[0, 256]) X.append(y) Y.append(piece == PIECE.EMPTY) classifier = make_pipeline(StandardScaler(), svm.SVC(C=10.0, gamma=0.01, probability=True)) classifier.fit(np.array(X), np.array(Y)) joblib.dump(classifier, f"classifiers/classifier_empty/white_piece_on_{square_color}_square.pkl") def train_pieces_svm() -> None: for piece in OUR_PIECES: for color in COLOR: total_weights = len(glob.glob(f"training_images/{piece}/{color}_square/*.png")) for piece in OUR_PIECES: for color in COLOR: current_weight = len(glob.glob(f"training_images/{piece}/{color}_square/*.png")) print(f"Training for piece: {piece}") X, Y = load_training_data(piece, color) classifier = svm.SVC(gamma=0.01, class_weight={0: current_weight, 1: total_weights}, probability=True) classifier.fit(X, Y) joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl") def train_pieces_svm_canny() -> None: for square_color in COLOR: X = [] Y = [] for piece in OUR_PIECES + (PIECE.EMPTY,): for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"): img = cv2.imread(filename) y, x = np.histogram(img.ravel(), bins=32, range=[0, 256]) X.append(y) Y.append(piece == PIECE.EMPTY) classifier = make_pipeline(StandardScaler(), svm.SVC(C=10.0, gamma=0.01, probability=True)) classifier.fit(np.array(X), np.array(Y)) joblib.dump(classifier, f"classifiers/classifier_empty/white_piece_on_{square_color}_square.pkl") def find_keypoints(camera_image: np.ndarray, baseline: np.ndarray, debug=False) -> Tuple[np.ndarray, np.ndarray]: """ Find keypoints in raw camera image of board. :return: (src points, dest points) """ cv2.imwrite("camera_image.png", camera_image) camera_image_gray = cv2.cvtColor(camera_image, cv2.COLOR_BGR2GRAY) #sift = cv2.xfeatures2d.SURF_create() kp_start = time.time() camera_image_keypoints = SIFT.detect(camera_image_gray, None) camera_image_keypoints, des = SIFT.compute(camera_image_gray, camera_image_keypoints) #print("kp:",time.time() - kp_start) def_flan = time.time() # 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_start = time.time() flann = cv2.FlannBasedMatcher(index_params, search_params) #print("end_def:", time.time() - def_flan) matches = flann.knnMatch(des, BASELINE_DES, k=2) #print("flann:",time.time() - flann_start) # Need to draw only good matches, so create a mask matchesMask = [[0, 0] for _ in range(len(matches))] # Ratio test as per Lowe's paper good_matches = [] for i, (m, n) in enumerate(matches): if m.distance < 0.55 * n.distance: matchesMask[i] = [1, 0] good_matches.append([m, n]) #good_matches = list(filter(lambda x: x[0].distance < 0.55 * x[1].distance, matches)) if debug: # Save keypoints keypoints_image = camera_image.copy() cv2.drawKeypoints(camera_image, keypoints=camera_image_keypoints, outImage=keypoints_image) cv2.imwrite("keypoints.png", keypoints_image) # Save matches matches_image = cv2.drawMatchesKnn( camera_image, camera_image_keypoints, baseline, BASELINE_KEYPOINTS, matches, None, matchColor=(0, 255, 0), singlePointColor=(255, 0, 0), matchesMask=matchesMask, flags=0 ) cv2.imwrite("matches.png", matches_image) # Extract location of good matches src_points = np.zeros((len(good_matches), 2), dtype=np.float32) dst_points = np.zeros((len(good_matches), 2), dtype=np.float32) for i, (m, n) in enumerate(good_matches): src_points[i, :] = camera_image_keypoints[m.queryIdx].pt dst_points[i, :] = BASELINE_KEYPOINTS[m.trainIdx].pt return src_points, dst_points def find_homography(camera_image: np.ndarray, baseline: np.ndarray = BASELINE, debug=False) -> np.ndarray: src_points, dst_points = find_keypoints(camera_image, baseline, debug=debug) h, mask = cv2.findHomography(src_points, dst_points, cv2.RANSAC) return h def warp_board(camera_image: np.ndarray, homography: np.ndarray = None, debug=False) -> np.ndarray: if homography is None: homography = find_homography(camera_image, BASELINE, debug=debug) height, width, channels = BASELINE.shape return cv2.warpPerspective(camera_image, homography, (width, height)) def get_square(warped_board: np.ndarray, position: POSITION) -> np.ndarray: width, _, _ = warped_board.shape # board is square anyway side = int(width * 0.045) size = width - 2 * side square_size = size // 8 padding = 2 x1 = side + (square_size * (position.file - 1)) x2 = x1 + square_size + padding y1 = max(0, side + (square_size * (8 - position.rank)) - padding) # 8 - rank because chessboard is from 8 to 1 y2 = min(width, y1 + square_size + padding) square = warped_board[y1:y2, x1:x2] return square def get_squares(warped_board: np.ndarray) -> Squares: # cv2.imwrite(f"warped_square_{square}.png", square) return {position: get_square(warped_board, position) for position in POSITION} def save_empty_fields(warped_board: np.ndarray, skip_rank: RANK = None, fourk=False) -> None: for position in POSITION: if position.rank == skip_rank: continue square = get_square(warped_board, position) if fourk: imwrite(f"training_images/4k/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square) else: imwrite(f"training_images/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square) def load_data_nn(spec_piece, color): X = None Y = None for piece in OUR_PIECES: piece_class = int(spec_piece == piece) for filename in glob.glob(f"training_images/{piece}/{color}_square/*.png"): image = cv2.imread(filename) 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: for _ in range(5): 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: for _ in range(5): X = np.vstack((X, data)) Y = np.vstack((Y, [piece_class])) return X, Y def train_nn(): for piece in OUR_PIECES: for color in COLOR: X, Y = load_data_nn(piece, color) classifier = neural_network.MLPClassifier(hidden_layer_sizes=256) classifier.fit(X, Y) joblib.dump(classifier, f"classifiers/neural_net_{piece}/{color}.pkl") if __name__ == '__main__': #train_nn() do_pre_processing() train_pieces_svm()