advancedskrald/main.py
2019-04-04 22:25:15 +02:00

622 lines
19 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
import sys
np.set_printoptions(threshold=sys.maxsize)
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_1554286488.605142_.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:
counter = 0
src = runner.get_square(warped, file, rank)
width, height, _ = src.shape
src = src[width//25:, height//25:]
# src = src[:-width//200, :-height//200]
segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.8, k=150, min_size=700)
segment = segmentator.processImage(src)
mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2)
masked = np.ma.masked_array(src, fill_value=0)
pls = []
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]
lel = (bottom - top) * (right - left)
#print(f"this is lel: {lel} ")
#print(f"this is meh: {np.sum(mask[:,:,0])} ")
if f"{file}{rank}" == "H7":
print("--"*20)
print("H7")
print(lel)
print(len(y))
print(np.max(segment))
# print(lel)
# print(np.sum(mask[:, :, 0]))
print("--"*20)
pls.append(len(y))
if len(y) < (164**2)*0.65:
counter += 1
cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{file}{rank}.jpg", dst)
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])):
print(f"{file}{rank} is nonempty")
non_empties.append([f"{file}{rank}", src])
print(counter)
print(np.max(segment))
empty += 1
print("++"*20)
print(counter)
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)