204 lines
7.4 KiB
Python
204 lines
7.4 KiB
Python
import cv2
|
|
import glob
|
|
import os
|
|
from datetime import datetime
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
from sklearn import cluster, metrics, svm
|
|
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
|
|
|
|
|
|
def generate_centers(number_of_clusters, sift: cv2.xfeatures2d_SIFT):
|
|
features = []
|
|
for piece in PIECE:
|
|
for color in COLOR:
|
|
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)
|
|
print(f"{piece}, {color}, {filename}")
|
|
features.append(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 PIECE:
|
|
for color in COLOR:
|
|
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(piece: PIECE, color: COLOR) -> Tuple[np.array, np.array]:
|
|
X = []
|
|
Y = []
|
|
for p in PIECE:
|
|
for filename in glob.glob(os.path.join("training_data", piece, f"{color}_square", "*.npy")):
|
|
data = np.load(filename)
|
|
X.append(data)
|
|
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 (PIECE.EMPTY, PIECE.ROOK, PIECE.KNIGHT):
|
|
for filename in glob.glob(os.path.join("training_images", f"{piece}", f"{square_color}_square", "*.png")):
|
|
img = cv2.imread(filename)
|
|
y, x = np.histogram(img.ravel(), bins=256, 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 PIECE:
|
|
for color in COLOR:
|
|
# 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 warp_board(camera_image, debug_image=None) -> np.ndarray:
|
|
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 _ 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])
|
|
|
|
img3 = 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.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
|
|
|
|
h, mask = cv2.findHomography(points1, points2, cv2.RANSAC)
|
|
|
|
height, width, channels = baseline.shape
|
|
return cv2.warpPerspective(camera_image, h, (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 = 0
|
|
|
|
x1 = side + (square_size * position.file)
|
|
x2 = x1 + square_size
|
|
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) -> None:
|
|
for position in POSITION:
|
|
if position.rank == skip_rank:
|
|
continue
|
|
square = get_square(warped_board, position)
|
|
imwrite(f"training_images/empty/{position.color}_square/training_{position}_{datetime.utcnow().timestamp()}.png", square)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
train_empty_or_piece_hist()
|