advancedskrald/main.py

177 lines
6.0 KiB
Python
Raw Normal View History

2019-04-04 20:25:15 +00:00
import sys
2019-04-10 11:39:50 +00:00
import warnings
from datetime import datetime
2019-04-10 11:39:50 +00:00
from typing import List, Tuple
2019-04-04 20:25:15 +00:00
2019-04-10 11:39:50 +00:00
import cv2
import matplotlib.pyplot as plt
import numpy as np
2019-04-10 11:39:50 +00:00
from sklearn.exceptions import DataConversionWarning
2019-04-04 10:59:37 +00:00
import runner
2019-04-10 11:39:50 +00:00
from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor, OUR_PIECES
2019-04-04 10:59:37 +00:00
warnings.filterwarnings(action='ignore', category=DataConversionWarning)
np.set_printoptions(threshold=sys.maxsize)
2019-04-04 10:59:37 +00:00
2019-04-10 11:39:50 +00:00
def identify_piece(image: np.ndarray, position: POSITION, sift: cv2.xfeatures2d_SIFT) -> PieceAndColor:
2019-04-04 10:59:37 +00:00
centers = np.load("training_data/centers.npy")
best = 0
2019-04-10 11:39:50 +00:00
probs = {p.name: {} for p in OUR_PIECES}
best_piece = best_color = None
2019-04-10 11:39:50 +00:00
for piece in OUR_PIECES:
for color in COLOR:
2019-04-04 10:59:37 +00:00
#color = runner.compute_color(file, rank)
classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl")
2019-04-04 10:59:37 +00:00
features = runner.generate_bag_of_words(image, centers, sift)
prob = classifier.predict_proba(features)
2019-04-10 11:39:50 +00:00
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:
2019-04-10 11:39:50 +00:00
best = prob[0, 1]
best_piece, best_color = piece, color
2019-04-10 11:39:50 +00:00
#print(probs)
return best_piece, best_color
2019-04-04 10:59:37 +00:00
def pred_test(position: POSITION, mystery_image=None, empty_bias=False):
2019-04-04 10:59:37 +00:00
sift = cv2.xfeatures2d.SIFT_create()
if mystery_image is None:
mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png")
2019-04-10 11:39:50 +00:00
probs = identify_piece(mystery_image, position, sift)
2019-04-04 10:59:37 +00:00
return probs
def pre_process_and_train() -> None:
2019-04-04 10:59:37 +00:00
runner.do_pre_processing()
runner.train_pieces_svm()
def build_board_from_squares(squares: Squares) -> Board:
2019-04-04 10:59:37 +00:00
sift = cv2.xfeatures2d.SIFT_create()
board = Board()
2019-04-04 10:59:37 +00:00
counter = 0
for position, square in squares.values():
2019-04-10 11:39:50 +00:00
likely_piece = identify_piece(square, position, sift)
board[position] = likely_piece
if likely_piece != PIECE.EMPTY:
2019-04-04 10:59:37 +00:00
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)
2019-04-04 10:59:37 +00:00
squares = runner.get_squares(warped)
board = build_board_from_squares(squares)
2019-04-04 10:59:37 +00:00
print(board)
2019-04-10 11:39:50 +00:00
def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])
2019-04-04 10:59:37 +00:00
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()
2019-04-04 10:59:37 +00:00
#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))
2019-04-10 11:39:50 +00:00
#print(f"{position}, {position.color}: {prob[0, 1]}")
if prob[0, 1] > 0.95:
2019-04-10 11:39:50 +00:00
# print(f"{position} is empty")
return PIECE.EMPTY
2019-04-04 10:59:37 +00:00
return None
2019-04-04 16:27:19 +00:00
def remove_most_empties(warped):
2019-04-04 16:27:19 +00:00
empty = 0
non_empties = []
for position in POSITION:
counter = 0
src = runner.get_square(warped, position)
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)
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"segment_test/segment_{datetime.utcnow().timestamp()}_{position}.png", dst)
2019-04-04 20:25:15 +00:00
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])):
2019-04-10 11:39:50 +00:00
#print(f"{position} is nonempty")
non_empties.append([position, src])
empty += 1
2019-04-04 16:27:19 +00:00
2019-04-10 11:39:50 +00:00
#print(64 - empty)
return non_empties
2019-04-10 11:39:50 +00:00
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:
2019-04-10 11:39:50 +00:00
if predict_empty(square, position) != PIECE.EMPTY:
completely_non_empties.append((position, square))
2019-04-04 10:59:37 +00:00
2019-04-10 11:39:50 +00:00
return completely_non_empties
2019-04-04 10:59:37 +00:00
2019-04-10 11:39:50 +00:00
if __name__ == '__main__':
2019-04-10 11:39:50 +00:00
#runner.train_pieces_svm()
2019-04-10 11:39:50 +00:00
#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)
2019-04-04 10:59:37 +00:00
2019-04-10 11:39:50 +00:00
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)
cv2.imshow("lal", rook_out)
cv2.imshow("lul", knight_out)
cv2.waitKey(0)
2019-04-04 10:59:37 +00:00
exit()
2019-04-10 11:39:50 +00:00
occupied = find_occupied_squares(warped)
2019-04-04 10:59:37 +00:00
2019-04-10 11:39:50 +00:00
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)
2019-04-04 10:59:37 +00:00
cv2.waitKey(0)