442 lines
14 KiB
Python
442 lines
14 KiB
Python
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()
|
|
|
|
|
|
|