588 lines
18 KiB
Python
588 lines
18 KiB
Python
from functools import lru_cache
|
|
|
|
import cv2
|
|
import runner
|
|
from sklearn.externals import joblib
|
|
import numpy as np
|
|
import operator
|
|
from matplotlib import pyplot as plt
|
|
import glob
|
|
import os
|
|
import heapq
|
|
import math
|
|
from datetime import datetime
|
|
|
|
pieces = ['rook', 'knight']
|
|
#pieces = ['rook', 'knight']
|
|
#piece_to_symbol = {'rook': 1, 'knight': 2, 'empty': 0}
|
|
piece_to_symbol = {'rook': 1, 'knight': 2}
|
|
colors = ['black', 'white']
|
|
|
|
|
|
|
|
def classify(image, sift : cv2.xfeatures2d_SIFT, file, rank, empty_bias=False):
|
|
centers = np.load("training_data/centers.npy")
|
|
probs = {'rook': {'black': 0, 'white': 0}, 'knight': {'black': 0, 'white': 0}, 'empty': {'black': 0, 'white': 0}}
|
|
#probs = {'rook': 0, 'knight': 0, 'empty': 0}
|
|
for piece in pieces:
|
|
for color in colors:
|
|
#color = runner.compute_color(file, rank)
|
|
classifier = joblib.load(f"classifiers/classifier_{piece}/{color}.pkl")
|
|
features = runner.generate_bag_of_words(image, centers, sift)
|
|
prob = classifier.predict_proba(features)
|
|
|
|
|
|
|
|
probs[piece][color] = prob[0, 1]
|
|
|
|
|
|
if empty_bias:
|
|
probs['empty'] *= 1.2
|
|
|
|
return probs
|
|
|
|
|
|
def pred_test(file, rank, mystery_image=None, empty_bias=False):
|
|
sift = cv2.xfeatures2d.SIFT_create()
|
|
if mystery_image is None:
|
|
mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png")
|
|
probs = classify(mystery_image, sift, file, rank, empty_bias=empty_bias)
|
|
return probs
|
|
|
|
|
|
|
|
def pre_process_and_train():
|
|
runner.do_pre_processing()
|
|
runner.train_pieces_svm()
|
|
|
|
|
|
def build_board_from_dict(board_dict : dict):
|
|
sift = cv2.xfeatures2d.SIFT_create()
|
|
board = [[0]*8 for _ in range(8)]
|
|
counter = 0
|
|
for idx, value in enumerate(board_dict.values()):
|
|
probs = classify(value, sift)
|
|
likely_piece = max(probs.items(), key=operator.itemgetter(1))[0]
|
|
symbol = piece_to_symbol[likely_piece]
|
|
|
|
column = idx // 8
|
|
row = (idx % 7)
|
|
|
|
|
|
board[row][column] = symbol
|
|
|
|
print(probs)
|
|
|
|
if likely_piece != 'empty':
|
|
counter += 1
|
|
|
|
print(counter)
|
|
print(64/(counter-1))
|
|
return board
|
|
|
|
|
|
|
|
|
|
def detect_using_nn(spec_image):
|
|
probs = {'rook': 0, 'knight': 0}
|
|
for piece in pieces:
|
|
piece_class = piece_to_symbol[piece]
|
|
win_size = (64, 64)
|
|
classifier = joblib.load("classifiers/neural_net_" + piece + ".pkl")
|
|
|
|
spec_image = cv2.resize(spec_image, (64, 128))
|
|
features = np.reshape(spec_image, (1, np.product(spec_image.shape)))
|
|
prob = classifier.predict_proba(features)
|
|
print(piece)
|
|
print(prob[0,1])
|
|
|
|
|
|
|
|
|
|
def test_entire_board():
|
|
board = cv2.imread("homo_pls_fuck.jpg")
|
|
warped = runner.warp_board(board)
|
|
|
|
board_dict = runner.get_squares(warped)
|
|
board = build_board_from_dict(board_dict)
|
|
print(board)
|
|
|
|
|
|
def lel_test():
|
|
# img = cv2.imread('training_images/rook/white/rook_training_D4_2.png')
|
|
|
|
counter = 0
|
|
|
|
for filename in glob.glob(os.path.join("training_images", "empty", "*", "*.png")):
|
|
|
|
img = cv2.imread(filename)
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
ret = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 3, 3)
|
|
|
|
# binarize the image
|
|
#ret, bw = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
|
|
# find connected components
|
|
connectivity = 4
|
|
nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(ret, connectivity, cv2.CV_32S)
|
|
sizes = stats[1:, -1]
|
|
nb_components = nb_components - 1
|
|
min_size = 250 # threshhold value for objects in scene
|
|
img2 = np.zeros((img.shape), np.uint8)
|
|
|
|
|
|
for i in range(0, nb_components + 1):
|
|
# use if sizes[i] >= min_size: to identify your objects
|
|
color = np.random.randint(255, size=3)
|
|
# draw the bounding rectangele around each object
|
|
cv2.rectangle(img2, (stats[i][0], stats[i][1]), (stats[i][0] + stats[i][2], stats[i][1] + stats[i][3]),
|
|
(0, 255, 0), 2)
|
|
img2[output == i + 1] = color
|
|
|
|
#print(nb_components+1)
|
|
|
|
if nb_components+1 >= 4:
|
|
counter += 1
|
|
print(filename)
|
|
cv2.imshow("lel", img2)
|
|
cv2.waitKey(0)
|
|
|
|
|
|
print(counter)
|
|
|
|
|
|
def selective_search(image, use_fast=False, use_slow=False):
|
|
# speed-up using multithreads
|
|
cv2.setUseOptimized(True)
|
|
cv2.setNumThreads(4)
|
|
|
|
|
|
if type(image) == str:
|
|
# read image
|
|
im = cv2.imread(image)
|
|
else:
|
|
im = image
|
|
# resize image
|
|
#newHeight = 200
|
|
#newWidth = int(im.shape[1] * 150 / im.shape[0])
|
|
#im = cv2.resize(im, (newWidth, newHeight))
|
|
|
|
#im = cv2.imread(image)
|
|
|
|
|
|
#lel, im = cv2.threshold(im, 128, 255, cv2.THRESH_BINARY)
|
|
|
|
|
|
|
|
# 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)
|
|
|
|
# Switch to fast but low recall Selective Search method
|
|
|
|
ss.switchToSingleStrategy()
|
|
|
|
if (use_fast):
|
|
ss.switchToSelectiveSearchFast()
|
|
|
|
# Switch to high recall but slow Selective Search method
|
|
elif (use_slow):
|
|
ss.switchToSelectiveSearchQuality()
|
|
|
|
# run selective search segmentation on input image
|
|
rects = ss.process()
|
|
#print('Total Number of Region Proposals: {}'.format(len(rects)))
|
|
|
|
# number of region proposals to show
|
|
numShowRects = 150
|
|
# increment to increase/decrease total number
|
|
# of reason proposals to be shown
|
|
increment = 1
|
|
|
|
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
|
|
# cv2.rectangle(imOut, (x, y), (x + w, y + h), (0, 255, 0), 1, cv2.LINE_AA)
|
|
|
|
# size = (max(w, x) - min(w, x)) * ((max(h, y) - min(h, y)))
|
|
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
|
|
#print(f"({x}, {y}), ({w}, {h})\n Of size: { size }")
|
|
|
|
#cv2.rectangle(imOut, (x, y), (x + w, y + h), (0, 255, 0), 1, cv2.LINE_AA)
|
|
#cv2.imshow("lel", imOut)
|
|
#cv2.waitKey(0)
|
|
|
|
best_proposals.append((rect, size))
|
|
|
|
#if size > biggest_size:
|
|
# biggest_rect = (x, y, w, h)
|
|
# biggest_size = size
|
|
# print(f"New biggest: \n({x}, {y}), ({w}, {h})\nOf size: {biggest_size}")
|
|
|
|
else:
|
|
break
|
|
|
|
height, width, channels = im.shape
|
|
center_x = width // 2
|
|
center_y = (height // 2)+5
|
|
dists = []
|
|
#print(f"Amount of best proposals:\n{len(best_proposals)}")
|
|
|
|
#print(f"lel: {len(heapq.nlargest(10, best_proposals, key=lambda x: x[1]))}")
|
|
for i in heapq.nlargest(10, best_proposals, key=lambda x: x[1]):
|
|
width, height, channels = im.shape
|
|
#print(width * height)
|
|
#print(i[1])
|
|
x, y, w, h = i[0]
|
|
|
|
|
|
|
|
if i[1] <= (width*height)*0.8 and i[1] > (width*height)*0.25:
|
|
|
|
|
|
imCop = imOut.copy()
|
|
|
|
#cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 2, cv2.LINE_AA)
|
|
#cv2.imshow("lel", imCop)
|
|
#cv2.waitKey(0)
|
|
|
|
|
|
#cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 4, cv2.LINE_AA)
|
|
|
|
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
|
|
|
|
#print(f"{box_center_x}, {box_center_y}, {center_x}, {center_y}")
|
|
|
|
dist = (center_x - box_center_x) ** 2 + (center_y - box_center_y) ** 2
|
|
print(dist)
|
|
dists.append([i, dist])
|
|
|
|
|
|
cv2.drawMarker(imCop, position=(x+w, h+y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x+w, y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x, y+h), color=(255, 0, 0), thickness=3)
|
|
|
|
cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3)
|
|
|
|
cv2.drawMarker(imCop, position=(center_x, center_y), color=(0, 0, 255), thickness=3)
|
|
|
|
|
|
#cv2.imshow("lel", imCop)
|
|
#cv2.waitKey(0)
|
|
|
|
#print("-------"*5)
|
|
|
|
for pls in dists:
|
|
imCop = imOut.copy()
|
|
x, y, w, h = pls[0][0]
|
|
#print(x,y,w,h)
|
|
#print(pls[1])
|
|
|
|
top_left = (x, y)
|
|
bottom_left = (x, y + h)
|
|
top_right = (x + w, y)
|
|
bottom_right = (x + w, y + h)
|
|
|
|
cv2.drawMarker(imCop, position=(x + w, h + y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x + w, y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x, y + h), color=(255, 0, 0), thickness=3)
|
|
|
|
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
|
|
|
|
cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3)
|
|
|
|
cv2.drawMarker(imCop, position=(center_y, center_x), color=(0, 0, 255), thickness=3)
|
|
|
|
|
|
cv2.rectangle(imCop, (x, y), (x + w, y + h), (0, 255, 0), 2, cv2.LINE_AA)
|
|
#cv2.imshow("lel", imCop)
|
|
#cv2.waitKey(0)
|
|
|
|
imCop = imOut.copy()
|
|
|
|
|
|
|
|
|
|
best = heapq.nsmallest(1, dists, key=lambda x: x[1])
|
|
|
|
if (len(best) == 0):
|
|
return ((0, 0), (0, height), (width, 0), (width, height))
|
|
|
|
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_left = (x, y + h)
|
|
top_right = (x + w, y)
|
|
bottom_right = (x + w, y + h)
|
|
|
|
cv2.drawMarker(imCop, position=(x + w, h + y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x + w, y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x, y), color=(255, 0, 0), thickness=3)
|
|
cv2.drawMarker(imCop, position=(x, y + h), color=(255, 0, 0), thickness=3)
|
|
|
|
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
|
|
|
|
cv2.drawMarker(imCop, position=(box_center_x, box_center_y), color=(0, 255, 0), thickness=3)
|
|
|
|
cv2.drawMarker(imCop, position=(center_x, center_y), color=(0, 0, 255), thickness=3)
|
|
|
|
|
|
#cv2.imshow("lel", imCop)
|
|
#cv2.waitKey(0)
|
|
return (top_left, bottom_left, top_right, bottom_right)
|
|
|
|
# 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 predict(square, file, rank):
|
|
square_color = runner.compute_color(file, rank)
|
|
|
|
y, x = np.histogram(square.ravel(), bins=256, range=[0, 256]) # TODO: Maybe img.ravel() ?
|
|
|
|
for color in ['black', 'white']:
|
|
empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{color}_square.pkl")
|
|
prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
|
|
print(f"{file}{rank}, {color}: {prob[0, 1]}")
|
|
if prob[0, 1] > 0.5:
|
|
return 'empty'
|
|
|
|
return None
|
|
|
|
|
|
|
|
@lru_cache()
|
|
def load_classifier(filename):
|
|
print(f"loading {filename}")
|
|
return joblib.load(filename)
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
board = cv2.imread("whole_boards/boards_for_empty/board_1554288951.972197_.png")
|
|
warped = runner.warp_board(board)
|
|
|
|
empty = 0
|
|
|
|
|
|
files = "ABCDEFGH"
|
|
ranks = [1, 2, 3, 4, 5, 6, 7, 8]
|
|
|
|
|
|
non_empties = []
|
|
|
|
for file in files:
|
|
for rank in ranks:
|
|
src = runner.get_square(warped, file, rank)
|
|
#src = cv2.GaussianBlur(src, (5, 5), 0)
|
|
segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.75, k=175, min_size=750)
|
|
segment = segmentator.processImage(src)
|
|
mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2)
|
|
masked = np.ma.masked_array(src, fill_value=0)
|
|
for i in range(np.max(segment)):
|
|
masked.mask = mask != i
|
|
y, x = np.where(segment == i)
|
|
top, bottom, left, right = min(y), max(y), min(x), max(x)
|
|
dst = masked.filled()[top: bottom + 1, left: right + 1]
|
|
cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{file}{rank}.jpg", dst)
|
|
|
|
if np.max(segment) >= 2:
|
|
print(f"{file}{rank} is nonempty")
|
|
non_empties.append([f"{file}{rank}", src])
|
|
print(np.max(segment))
|
|
empty += 1
|
|
print(64-empty)
|
|
|
|
for non_empty in non_empties:
|
|
cv2.imshow(non_empty[0], non_empty[1])
|
|
cv2.waitKey(0)
|
|
exit()
|
|
|
|
#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()
|
|
|
|
|
|
board = cv2.imread("whole_boards/board_102_1554110461.608167_.png")
|
|
warped = runner.warp_board(board)
|
|
|
|
files = "ABCDEFGH"
|
|
ranks = [1, 2, 3, 4, 5, 6, 7, 8]
|
|
|
|
counter = 0
|
|
|
|
for file in files:
|
|
for rank in ranks:
|
|
square = runner.get_square(warped, file, rank)
|
|
if predict(square, file, rank) == 'empty':
|
|
counter += 1
|
|
|
|
print(counter)
|
|
exit()
|
|
|
|
|
|
square = runner.get_square(warped, "D", 2)
|
|
|
|
gray_square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
|
|
print(cv2.meanStdDev(gray_square)[1])
|
|
print(cv2.meanStdDev(square)[1])
|
|
cv2.imshow("square", square)
|
|
cv2.waitKey(0)
|
|
|
|
|
|
|
|
|
|
print(pred_test("C", 2, square))
|
|
|
|
sift: cv2.xfeatures2d_SIFT = cv2.xfeatures2d.SIFT_create()
|
|
gray = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
|
|
|
|
kp, desc = sift.detectAndCompute(gray, None)
|
|
|
|
cv2.drawKeypoints(square, kp, square)
|
|
|
|
cv2.imshow("kp", square)
|
|
cv2.waitKey(0)
|
|
|
|
exit()
|
|
|
|
board = cv2.imread("whole_boards/board_202_1554154094.001122_.png")
|
|
|
|
runner.fetch_empty_fields(board)
|
|
exit()
|
|
|
|
|
|
warped = runner.warp_board(board)
|
|
|
|
counter = 0
|
|
|
|
#square = runner.get_square(warped, "A", 3)
|
|
|
|
#top_left, bottom_left, top_right, bottom_right = selective_search(square, use_fast=True)
|
|
#cropped = square[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
|
|
|
|
|
|
|
|
for file in files:
|
|
for rank in ranks:
|
|
|
|
square = runner.get_square(warped, file, rank)
|
|
|
|
|
|
top_left, bottom_left, top_right, bottom_right = selective_search(square, use_fast=True)
|
|
cropped = square[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
|
|
|
|
rect_width = bottom_right[0] - bottom_left[0]
|
|
rect_height = bottom_right[1] - top_right[1]
|
|
|
|
size = rect_width * rect_height
|
|
|
|
square_height, square_width, channels = square.shape
|
|
|
|
empty_bias = (size == square_height*square_width)
|
|
if size == square_height*square_width:
|
|
print(f"{file}{rank} is likely empty")
|
|
|
|
res = pred_test(file, rank, mystery_image=square, empty_bias=empty_bias)
|
|
print(res)
|
|
if (max(res.items(), key=operator.itemgetter(1))[0] == 'empty'):
|
|
counter += 1
|
|
|
|
|
|
|
|
print(f"Amount of empty fields: {counter}")
|
|
#print("Non-cropped:\t",pred_test(square))
|
|
#print("Cropped:\t",pred_test(cropped))
|
|
|
|
#cv2.imshow("square", square)
|
|
#cv2.waitKey(0)
|
|
|
|
|
|
#runner.do_pre_processing()
|
|
#runner.train()
|
|
|
|
|
|
#img = "warped_square_B5.png"
|
|
#detect_using_nn(img)
|
|
|
|
#selective_search("training_images/empty/colorless/warped_square_A6.png", use_fast=True)
|
|
#selective_search("warped_square_B5.png", use_fast=True)
|
|
img = "training_images/rook/white/rook_training_D4_7.png"
|
|
#img = "training_images/rook/white_square/rook_training_E4_10.png"
|
|
#img = "training_images/knight/white_square/training_D5_134.png"
|
|
|
|
#top_left, bottom_left, top_right, bottom_right = selective_search(img, use_fast=True)
|
|
#cropped = cv2.imread(img)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
|
|
|
|
#cv2.imshow("output", cropped)
|
|
#print(pred_test(cropped))
|
|
#cv2.waitKey(0)
|
|
|
|
|
|
|
|
#lel_test()
|
|
|
|
|
|
# test_entire_board()
|
|
|
|
#board = [[0, 0, 1, 2, 0, 0, 0, 2], [0, 1, 2, 2, 1, 0, 0, 1], [0, 0, 0, 0, 1, 0, 2, 0], [0, 2, 2, 1, 1, 2, 2, 0], [0, 1, 0, 0, 1, 2, 0, 0], [0, 0, 0, 0, 0, 2, 2, 0], [0, 0, 0, 2, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0]]
|
|
#for i in board:
|
|
# print(i)
|
|
|
|
#warped = cv2.imread("homo_pls_fuck.jpg")
|
|
#square = runner.get_square(warped, "D", 4)
|
|
#print(pred_test(square))
|
|
#cv2.imshow("lel", square)
|
|
#cv2.waitKey(0)
|