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 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
|
||||
|
|
Loading…
Reference in New Issue
Block a user