Maybe faster maybe not.
This commit is contained in:
parent
cde518c14e
commit
241c434672
16
main.py
16
main.py
|
@ -1,7 +1,5 @@
|
||||||
import sys
|
import sys
|
||||||
import warnings
|
import warnings
|
||||||
from datetime import datetime
|
|
||||||
from typing import List, Tuple
|
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
|
@ -80,7 +78,7 @@ def test_entire_board() -> None:
|
||||||
print(board)
|
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])
|
y, x = np.histogram(square.ravel(), bins=32, range=[0, 256])
|
||||||
|
|
||||||
left, right = x[:-1], x[1:]
|
left, right = x[:-1], x[1:]
|
||||||
|
@ -91,16 +89,10 @@ def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
|
||||||
plt.xlabel(f"{position}")
|
plt.xlabel(f"{position}")
|
||||||
#plt.show()
|
#plt.show()
|
||||||
|
|
||||||
|
|
||||||
#for color in COLOR:
|
|
||||||
empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{position.color}_square.pkl")
|
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))
|
prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
|
||||||
#print(f"{position}, {position.color}: {prob[0, 1]}")
|
#print(f"{position}, {position.color}: {prob[0, 1]}")
|
||||||
if prob[0, 1] > 0.95:
|
return prob[0, 1] > 0.95
|
||||||
# print(f"{position} is empty")
|
|
||||||
return PIECE.EMPTY
|
|
||||||
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def remove_most_empties(warped):
|
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)
|
top, bottom, left, right = min(y), max(y), min(x), max(x)
|
||||||
|
|
||||||
dst = masked.filled()[top: bottom + 1, left: right + 1]
|
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 (
|
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])):
|
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 = {}
|
completely_non_empties = {}
|
||||||
for position, square in 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
|
completely_non_empties[position] = square
|
||||||
|
|
||||||
return completely_non_empties
|
return completely_non_empties
|
||||||
|
|
Loading…
Reference in New Issue
Block a user