Maybe faster maybe not.

This commit is contained in:
Casper 2019-04-24 09:54:37 +02:00
parent cde518c14e
commit 241c434672

16
main.py
View File

@ -1,7 +1,5 @@
import sys
import warnings
from datetime import datetime
from typing import List, Tuple
import cv2
import matplotlib.pyplot as plt
@ -80,7 +78,7 @@ def test_entire_board() -> None:
print(board)
def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
def predict_empty(square: np.ndarray, position: POSITION) -> bool:
y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])
left, right = x[:-1], x[1:]
@ -91,16 +89,10 @@ def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
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
return prob[0, 1] > 0.95
def remove_most_empties(warped):
@ -128,7 +120,7 @@ def remove_most_empties(warped):
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)
#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])):
@ -146,7 +138,7 @@ def find_occupied_squares(warped: np.ndarray) -> Squares:
completely_non_empties = {}
for position, square in non_empties:
if predict_empty(square, position) != PIECE.EMPTY:
if not predict_empty(square, position):
completely_non_empties[position] = square
return completely_non_empties