Works a bit.

This commit is contained in:
Alexander Munch-Hansen 2019-04-10 13:39:50 +02:00
parent 6393b12e82
commit 3e4327daee
3 changed files with 78 additions and 95 deletions

123
main.py
View File

@ -1,39 +1,39 @@
import cv2
import sys import sys
from collections import defaultdict import warnings
from datetime import datetime from datetime import datetime
from typing import List, Tuple
import cv2
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import numpy as np
from sklearn.exceptions import DataConversionWarning
import runner import runner
from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor, OUR_PIECES
from sklearn.exceptions import DataConversionWarning
import warnings
warnings.filterwarnings(action='ignore', category=DataConversionWarning) warnings.filterwarnings(action='ignore', category=DataConversionWarning)
np.set_printoptions(threshold=sys.maxsize) np.set_printoptions(threshold=sys.maxsize)
def identify_piece(image: np.ndarray, sift : cv2.xfeatures2d_SIFT, empty_bias=False) -> PieceAndColor: def identify_piece(image: np.ndarray, position: POSITION, sift: cv2.xfeatures2d_SIFT) -> PieceAndColor:
centers = np.load("training_data/centers.npy") centers = np.load("training_data/centers.npy")
probs = defaultdict(lambda: defaultdict(float))
best = 0 best = 0
probs = {p.name: {} for p in OUR_PIECES}
best_piece = best_color = None best_piece = best_color = None
for piece in PIECE: for piece in OUR_PIECES:
for color in COLOR: for color in COLOR:
#color = runner.compute_color(file, rank) #color = runner.compute_color(file, rank)
classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl") classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl")
features = runner.generate_bag_of_words(image, centers, sift) features = runner.generate_bag_of_words(image, centers, sift)
prob = classifier.predict_proba(features) prob = classifier.predict_proba(features)
probs[piece][color] = prob[0, 1] probs[piece.name][color.name] = prob[0, 1]
print(f"{piece}, {color}, {prob[0, 1]}")
#if prob[0, 1] > best and color == position.color: # can only be best if correct color. Iterating through both colors for debugging only
if prob[0, 1] > best: if prob[0, 1] > best:
best = prob[0, 1]
best_piece, best_color = piece, color best_piece, best_color = piece, color
print(probs) #print(probs)
if empty_bias:
probs[PIECE.EMPTY] *= 1.2
return best_piece, best_color return best_piece, best_color
@ -41,7 +41,7 @@ def pred_test(position: POSITION, mystery_image=None, empty_bias=False):
sift = cv2.xfeatures2d.SIFT_create() sift = cv2.xfeatures2d.SIFT_create()
if mystery_image is None: if mystery_image is None:
mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png") mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png")
probs = identify_piece(mystery_image, sift, empty_bias=empty_bias) probs = identify_piece(mystery_image, position, sift)
return probs return probs
@ -55,7 +55,7 @@ def build_board_from_squares(squares: Squares) -> Board:
board = Board() board = Board()
counter = 0 counter = 0
for position, square in squares.values(): for position, square in squares.values():
likely_piece = identify_piece(square, sift) likely_piece = identify_piece(square, position, sift)
board[position] = likely_piece board[position] = likely_piece
if likely_piece != PIECE.EMPTY: if likely_piece != PIECE.EMPTY:
counter += 1 counter += 1
@ -74,7 +74,7 @@ def test_entire_board() -> None:
print(board) print(board)
def predict(square: np.ndarray, position: POSITION) -> PIECE: def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
y, x = np.histogram(square.ravel(), bins=32, range=[0, 256]) y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])
left, right = x[:-1], x[1:] left, right = x[:-1], x[1:]
@ -89,9 +89,9 @@ def predict(square: np.ndarray, position: POSITION) -> PIECE:
#for color in COLOR: #for color in COLOR:
empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{position.color}_square.pkl") empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{position.color}_square.pkl")
prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1)) prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
print(f"{position}, {position.color}: {prob[0, 1]}") #print(f"{position}, {position.color}: {prob[0, 1]}")
if prob[0, 1] > 0.95: if prob[0, 1] > 0.95:
print(f"{position} is empty") # print(f"{position} is empty")
return PIECE.EMPTY return PIECE.EMPTY
return None return None
@ -124,82 +124,53 @@ def remove_most_empties(warped):
if np.max(segment) > 0 and not np.all([x < (164 ** 2) * 0.2 for x in pls]) and ( if np.max(segment) > 0 and not np.all([x < (164 ** 2) * 0.2 for x in pls]) and (
np.max(segment) >= 3 or np.all([x < (164 ** 2) * 0.942 for x in pls])): np.max(segment) >= 3 or np.all([x < (164 ** 2) * 0.942 for x in pls])):
print(f"{position} is nonempty") #print(f"{position} is nonempty")
non_empties.append([position, src]) non_empties.append([position, src])
empty += 1 empty += 1
print(64 - empty) #print(64 - empty)
return non_empties return non_empties
if __name__ == '__main__':
#board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
board = cv2.imread("whole_boards/boards_for_empty/board_1554288606.075646_rank_1.png")
warped = runner.warp_board(board)
def find_occupied_squares(warped: np.ndarray) -> List[Tuple[POSITION, np.ndarray]]:
non_empties = remove_most_empties(warped) non_empties = remove_most_empties(warped)
#empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_white_square.pkl")
#print(empty_classifier.predict_proba(np.array([0]*16).reshape(1, -1))[0, 1])
#exit()
counter = 0
completely_non_empties = [] completely_non_empties = []
for position, square in non_empties: for position, square in non_empties:
#predict(square, position) if predict_empty(square, position) != PIECE.EMPTY:
completely_non_empties.append((position, square))
#y, x = np.histogram(square.ravel(), bins=32, range=[0, 256]) return completely_non_empties
#left, right = x[:-1], x[1:]
#X = np.array([left, right]).T.flatten()
#Y = np.array([y, y]).T.flatten()
#plt.plot(X, Y)
#plt.xlabel(f"{position}")
#plt.show()
if __name__ == '__main__':
if predict(square,position) == PIECE.EMPTY: #runner.train_pieces_svm()
counter += 1
else:
completely_non_empties.append([position, square])
#board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
board = cv2.imread("whole_boards/boards_for_empty/board_1554286515.323962_rank_3.png")
warped = runner.warp_board(board)
rook_square = runner.get_square(warped, POSITION.H3)
knight_square = runner.get_square(warped, POSITION.D3)
print(counter) cv2.imshow("lel", rook_square)
for position, square in completely_non_empties: cv2.imshow("lil", knight_square)
cv2.imshow(f"{position}", square) rook_out = cv2.Canny(rook_square, 50, 55, L2gradient=True)
knight_out = cv2.Canny(knight_square, 50, 55, L2gradient=True)
cv2.imshow("lal", rook_out)
cv2.imshow("lul", knight_out)
cv2.waitKey(0) cv2.waitKey(0)
exit() exit()
occupied = find_occupied_squares(warped)
square_img = runner.get_square(warped, "D", 2) sift = cv2.xfeatures2d.SIFT_create()
for position, square in occupied:
gray_square_img = cv2.cvtColor(square_img, cv2.COLOR_BGR2GRAY) print("---"*15)
print(cv2.meanStdDev(gray_square_img)[1]) piece, color = identify_piece(square, position, sift)
print(cv2.meanStdDev(square_img)[1]) print(f"{piece} on {position}")
cv2.imshow("square", square_img) text_color = 255 if color == COLOR.WHITE else 0
cv2.putText(square, f"{position} {piece.name}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(text_color,)*3, thickness=3)
cv2.imshow(f"{position}", square)
cv2.waitKey(0) cv2.waitKey(0)
print(pred_test("C", 2, square_img))
sift: cv2.xfeatures2d_SIFT = cv2.xfeatures2d.SIFT_create()
gray = cv2.cvtColor(square_img, cv2.COLOR_BGR2GRAY)
kp, desc = sift.detectAndCompute(gray, None)
cv2.drawKeypoints(square_img, kp, square_img)
cv2.imshow("kp", square_img)
cv2.waitKey(0)

View File

@ -10,20 +10,26 @@ from sklearn.externals import joblib
from sklearn.pipeline import make_pipeline from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler from sklearn.preprocessing import StandardScaler
from util import RANK, POSITION, imwrite, PIECE, COLOR, Squares from util import RANK, POSITION, imwrite, PIECE, COLOR, Squares, OUR_PIECES
def generate_centers(number_of_clusters, sift: cv2.xfeatures2d_SIFT): def generate_centers(number_of_clusters, sift: cv2.xfeatures2d_SIFT):
features = [] features = None
for piece in PIECE: for piece in OUR_PIECES:
for color in COLOR: for color in COLOR:
for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")): for filename in glob.glob(f"training_images/{piece}/{color}_square/*.png"):
image = cv2.imread(filename) image = cv2.imread(filename)
#image = selective_search(image, use_fast=True) #image = selective_search(image, use_fast=True)
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
kp, desc = sift.detectAndCompute(gray, None) kp, desc = sift.detectAndCompute(gray, None)
print(f"{piece}, {color}, {filename}") print(f"{piece}, {color}, {filename}")
features.append(desc)
if features is None:
features = np.array(desc)
else:
print(f"{piece}, {color}, {filename}")
features = np.vstack((features, desc))
features = np.array(features) features = np.array(features)
k_means = cluster.KMeans(number_of_clusters) k_means = cluster.KMeans(number_of_clusters)
k_means.fit(features) k_means.fit(features)
@ -55,9 +61,9 @@ def do_pre_processing() -> None:
np.save("training_data/centers", centers) np.save("training_data/centers", centers)
for piece in PIECE: for piece in OUR_PIECES:
for color in COLOR: for color in COLOR:
for filename in glob.glob(os.path.join("training_images", piece, f"{color}_square", "*.png")): for filename in glob.glob(f"training_images/{piece}/{color}_square/*.png"):
image = cv2.imread(filename) image = cv2.imread(filename)
#image = selective_search(image, image_name=filename, use_fast=True) #image = selective_search(image, image_name=filename, use_fast=True)
bow_features = generate_bag_of_words(image, centers, sift) bow_features = generate_bag_of_words(image, centers, sift)
@ -67,10 +73,10 @@ def do_pre_processing() -> None:
def load_training_data(piece: PIECE, color: COLOR) -> Tuple[np.array, np.array]: def load_training_data(piece: PIECE, color: COLOR) -> Tuple[np.array, np.array]:
X = [] X = []
Y = [] Y = []
for p in PIECE: for p in OUR_PIECES:
for filename in glob.glob(os.path.join("training_data", piece, f"{color}_square", "*.npy")): for filename in glob.glob(f"training_data/{piece}/{color}_square/*.npy"):
data = np.load(filename) data = np.load(filename)
X.append(data) X.append(data[0])
Y.append(p == piece) Y.append(p == piece)
return np.array(X), np.array(Y) return np.array(X), np.array(Y)
@ -79,8 +85,8 @@ def train_empty_or_piece_hist() -> None:
for square_color in COLOR: for square_color in COLOR:
X = [] X = []
Y = [] Y = []
for piece in (PIECE.EMPTY, PIECE.ROOK, PIECE.KNIGHT): for piece in OUR_PIECES + (PIECE.EMPTY,):
for filename in glob.glob(os.path.join("training_images", f"{piece}", f"{square_color}_square", "*.png")): for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"):
img = cv2.imread(filename) img = cv2.imread(filename)
y, x = np.histogram(img.ravel(), bins=32, range=[0, 256]) y, x = np.histogram(img.ravel(), bins=32, range=[0, 256])
X.append(y) X.append(y)
@ -93,15 +99,16 @@ def train_empty_or_piece_hist() -> None:
def train_pieces_svm() -> None: def train_pieces_svm() -> None:
for piece in PIECE: for piece in OUR_PIECES:
for color in COLOR: for color in COLOR:
# TODO: Consider removing empty from total_weights, so all classifiers do not consider empty pieces total_weights = len(glob.glob(f"training_images/{piece}/{color}_square/*.png"))
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}") 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) X, Y = load_training_data(piece, color)
classifier = svm.SVC(class_weight={0: current_weight, 1: total_weights - current_weight}, probability=True) classifier = svm.SVC(C=10, gamma=0.01, class_weight={0: 45, 1: 1}, probability=True)
classifier.fit(X, Y) classifier.fit(X, Y)
joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl") joblib.dump(classifier, f"classifiers/classifier_{piece}/{color}.pkl")

View File

@ -33,6 +33,11 @@ class PIECE(Enum):
PieceAndColor = Tuple[PIECE, COLOR] PieceAndColor = Tuple[PIECE, COLOR]
OUR_PIECES = (
PIECE.ROOK,
PIECE.KNIGHT,
)
class FILE(int, Enum): class FILE(int, Enum):
A = 1 A = 1
@ -88,5 +93,5 @@ def imwrite(*args, **kwargs):
@lru_cache() @lru_cache()
def load_classifier(filename): def load_classifier(filename):
print(f"Loading classifier {filename}") # print(f"Loading classifier {filename}")
return joblib.load(filename) return joblib.load(filename)