198 lines
6.6 KiB
Python
198 lines
6.6 KiB
Python
import sys
|
|
import warnings
|
|
from datetime import datetime
|
|
from typing import List, Tuple
|
|
|
|
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from sklearn.exceptions import DataConversionWarning
|
|
|
|
import runner
|
|
from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor, OUR_PIECES
|
|
|
|
warnings.filterwarnings(action='ignore', category=DataConversionWarning)
|
|
np.set_printoptions(threshold=sys.maxsize)
|
|
|
|
|
|
def identify_piece(image: np.ndarray, position: POSITION, sift: cv2.xfeatures2d_SIFT) -> PieceAndColor:
|
|
centers = np.load("training_data/centers.npy")
|
|
best = 0
|
|
probs = {p.name: {} for p in OUR_PIECES}
|
|
|
|
best_piece = best_color = None
|
|
for piece in OUR_PIECES:
|
|
for color in COLOR:
|
|
#color = runner.compute_color(file, rank)
|
|
#classifier = load_classifier(f"classifiers/neural_net_{piece}/{color}.pkl")
|
|
classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl")
|
|
features = runner.generate_bag_of_words(image, centers, sift)
|
|
prob = classifier.predict_proba(features)
|
|
|
|
#image = cv2.resize(image, (172, 172))
|
|
#data = np.reshape(image, (1, np.product(image.shape)))
|
|
#prob = classifier.predict_proba(data)
|
|
|
|
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:
|
|
best = prob[0, 1]
|
|
best_piece, best_color = piece, color
|
|
#print(probs)
|
|
return best_piece, best_color
|
|
|
|
|
|
def pred_test(position: POSITION, 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 = identify_piece(mystery_image, position, sift)
|
|
return probs
|
|
|
|
|
|
def pre_process_and_train() -> None:
|
|
runner.do_pre_processing()
|
|
runner.train_pieces_svm()
|
|
|
|
|
|
def build_board_from_squares(squares: Squares) -> Board:
|
|
sift = cv2.xfeatures2d.SIFT_create()
|
|
board = Board()
|
|
counter = 0
|
|
for position, square in squares.values():
|
|
likely_piece = identify_piece(square, position, sift)
|
|
board[position] = likely_piece
|
|
if likely_piece != PIECE.EMPTY:
|
|
counter += 1
|
|
|
|
print(counter)
|
|
print(64/(counter-1))
|
|
return board
|
|
|
|
|
|
def test_entire_board() -> None:
|
|
board_img = cv2.imread("homo_pls_fuck.jpg")
|
|
warped = runner.warp_board(board_img)
|
|
|
|
squares = runner.get_squares(warped)
|
|
board = build_board_from_squares(squares)
|
|
print(board)
|
|
|
|
|
|
def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
|
|
y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])
|
|
|
|
left, right = x[:-1], x[1:]
|
|
X = np.array([left, right]).T.flatten()
|
|
Y = np.array([y, y]).T.flatten()
|
|
area = sum(np.diff(x) * y)
|
|
plt.plot(X, Y)
|
|
plt.xlabel(f"{position}")
|
|
#plt.show()
|
|
|
|
|
|
#for color in COLOR:
|
|
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))
|
|
#print(f"{position}, {position.color}: {prob[0, 1]}")
|
|
if prob[0, 1] > 0.95:
|
|
# print(f"{position} is empty")
|
|
return PIECE.EMPTY
|
|
|
|
return None
|
|
|
|
def remove_most_empties(warped):
|
|
empty = 0
|
|
non_empties = []
|
|
for position in POSITION:
|
|
counter = 0
|
|
|
|
img_src = runner.get_square(warped, position)
|
|
width, height, _ = img_src.shape
|
|
src = img_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)
|
|
|
|
|
|
if position == POSITION.E8:
|
|
print(np.max(segment))
|
|
print(len(y))
|
|
|
|
|
|
pls.append(len(y))
|
|
top, bottom, left, right = min(y), max(y), min(x), max(x)
|
|
|
|
dst = masked.filled()[top: bottom + 1, left: right + 1]
|
|
cv2.imwrite(f"tmp_seg/segment_{datetime.utcnow().timestamp()}_{position}.png", 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.9469 for x in pls])):
|
|
#print(f"{position} is nonempty")
|
|
non_empties.append([position, img_src])
|
|
empty += 1
|
|
|
|
#print(64 - empty)
|
|
|
|
return non_empties
|
|
|
|
|
|
def find_occupied_squares(warped: np.ndarray) -> List[Tuple[POSITION, np.ndarray]]:
|
|
non_empties = remove_most_empties(warped)
|
|
|
|
completely_non_empties = []
|
|
for position, square in non_empties:
|
|
if predict_empty(square, position) != PIECE.EMPTY:
|
|
completely_non_empties.append((position, square))
|
|
|
|
return completely_non_empties
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
#runner.train_pieces_svm()
|
|
|
|
board = cv2.imread("whole_boards/boards_for_empty/board_1554288891.129901_rank_8.png")
|
|
#board = cv2.imread("whole_boards/boards_for_empty/board_1554286515.323962_rank_3.png")
|
|
warped = runner.warp_board(board)
|
|
|
|
tmp = find_occupied_squares(warped)
|
|
for pos, square in tmp:
|
|
cv2.imshow(f"{pos}", square)
|
|
cv2.waitKey(0)
|
|
exit()
|
|
|
|
"""
|
|
rook_square = runner.get_square(warped, POSITION.H3)
|
|
knight_square = runner.get_square(warped, POSITION.D3)
|
|
cv2.imshow("lel", rook_square)
|
|
cv2.imshow("lil", knight_square)
|
|
#rook_out = cv2.Canny(rook_square, 50, 55, L2gradient=True)
|
|
knight_out = cv2.Canny(knight_square, 50, 55, L2gradient=True)
|
|
knight_out_l = cv2.Canny(knight_square, 50, 55, L2gradient=False)
|
|
cv2.imshow("lal", knight_out)
|
|
cv2.imshow("lul", knight_out_l)
|
|
cv2.waitKey(0)
|
|
exit()
|
|
"""
|
|
occupied = find_occupied_squares(warped)
|
|
|
|
sift = cv2.xfeatures2d.SIFT_create()
|
|
for position, square in occupied:
|
|
print("---"*15)
|
|
piece, color = identify_piece(square, position, sift)
|
|
print(f"{piece} on {position}")
|
|
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)
|