Now we works and we still have type hints so thats nice

This commit is contained in:
Casper 2019-04-17 19:58:03 +02:00
parent 0f702640a3
commit 417538d768
10 changed files with 75 additions and 124 deletions

View File

@ -1,6 +1,6 @@
## Test Adapter:
```bash
base64 --wrap=0 < whole_boards/board_118_1554110522.620303_.png| python3 adapter.py
base64 --wrap=0 < whole_boards/boards_for_empty/board_1554285167.655788_rank_5.png | python3 adapter.py
```
## Build OpenCV

View File

@ -5,21 +5,29 @@ import sys
import cv2
import numpy as np
from runner import find_homography
from main import find_occupied_squares
from runner import find_homography, warp_board
# Load base64 encoded image from stdin
from tensor_classifier import predict_board
stdin = sys.stdin.readline()
stdin_decoded = base64.b64decode(stdin)
img_array = np.frombuffer(stdin_decoded, dtype=np.uint8)
camera_img = cv2.imdecode(img_array, flags=cv2.COLOR_BGR2RGB)
camera_img = cv2.cvtColor(camera_img, cv2.COLOR_BGR2RGB)
# Find keypoints in image and pass them back to unity
# def do_everything:
homography = find_homography(camera_img)
warped_board = warp_board(camera_img, homography)
occupied_squares = find_occupied_squares(warped_board)
board = predict_board(occupied_squares)
# Finally, output to stdout for unity to read
result = {
"homography": homography.tolist(),
"board": board.to_array,
}
print(json.dumps(result))

13
main.py
View File

@ -102,6 +102,7 @@ def predict_empty(square: np.ndarray, position: POSITION) -> PIECE:
return None
def remove_most_empties(warped):
empty = 0
non_empties = []
@ -123,12 +124,6 @@ def remove_most_empties(warped):
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)
@ -146,13 +141,13 @@ def remove_most_empties(warped):
return non_empties
def find_occupied_squares(warped: np.ndarray) -> List[Tuple[POSITION, np.ndarray]]:
def find_occupied_squares(warped: np.ndarray) -> Squares:
non_empties = remove_most_empties(warped)
completely_non_empties = []
completely_non_empties = {}
for position, square in non_empties:
if predict_empty(square, position) != PIECE.EMPTY:
completely_non_empties.append((position, square))
completely_non_empties[position] = square
return completely_non_empties

View File

@ -10,11 +10,11 @@ cap = cv2.VideoCapture(0)
color = COLOR.BLACK
rank = RANK.EIGHT
pieces = {
PIECE.rook: [POSITION(FILE.A, rank), POSITION(FILE.F, rank)],
PIECE.knight: [POSITION(FILE.E, rank), POSITION(FILE.H, rank)],
PIECE.bishop: [POSITION(FILE.C, rank), POSITION(FILE.D, rank)],
PIECE.queen: [POSITION(FILE.B, rank)],
PIECE.king: [POSITION(FILE.G, rank)],
PIECE.rook: [POSITION((FILE.A, rank)), POSITION((FILE.F, rank))],
PIECE.knight: [POSITION((FILE.E, rank)), POSITION((FILE.H, rank))],
PIECE.bishop: [POSITION((FILE.C, rank)), POSITION((FILE.D, rank))],
PIECE.queen: [POSITION((FILE.B, rank))],
PIECE.king: [POSITION((FILE.G, rank))],
}
while True:

View File

@ -89,6 +89,7 @@ def train_empty_or_piece_hist() -> None:
X = []
Y = []
for piece in OUR_PIECES + (PIECE.EMPTY,):
print(f"training training_images/{piece}/{square_color}_square/*.png")
for filename in glob.glob(f"training_images/{piece}/{square_color}_square/*.png"):
img = cv2.imread(filename)
y, x = np.histogram(img.ravel(), bins=32, range=[0, 256])
@ -209,7 +210,8 @@ def find_homography(camera_image: np.ndarray,
def warp_board(camera_image: np.ndarray, homography: np.ndarray = None, debug=False) -> np.ndarray:
baseline = cv2.imread(str(here.joinpath("new_baseline_board.png")))
homography = homography or find_homography(camera_image, baseline, debug=debug)
if homography is None:
homography = find_homography(camera_image, baseline, debug=debug)
height, width, channels = baseline.shape
return cv2.warpPerspective(camera_image, homography, (width, height))
@ -289,4 +291,5 @@ def train_nn():
if __name__ == '__main__':
train_nn()
#train_nn()
train_empty_or_piece_hist()

View File

@ -1,71 +1,27 @@
#from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
from tensorflow.python.keras import datasets, layers, models
import glob
import numpy as np
import cv2
import numpy as np
from tensorflow.python.keras import models
import runner
from main import find_occupied_squares
from util import POSITION
from util import PIECE, Squares, Board
new_model = models.load_model('chess_model_3_pieces.h5')
new_model.summary()
#new_model.summary()
#board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
#board = cv2.imread("whole_boards/boards_for_empty/board_1554285167.655788_rank_5.png")
board = cv2.imread("whole_boards/boards_for_empty/board_1554288891.129901_rank_8.png")
warped = runner.warp_board(board)
occupied = find_occupied_squares(warped)
pos_1 = POSITION.C5
pos_2 = POSITION.D5
pos_3 = POSITION.G5
pos_4 = POSITION.H5
square_1 = runner.get_square(warped, pos_1)
square_2 = runner.get_square(warped, pos_2)
square_3 = runner.get_square(warped, pos_3)
square_4 = runner.get_square(warped, pos_4)
square_1 = cv2.cvtColor(square_1, cv2.COLOR_BGR2GRAY)
square_2 = cv2.cvtColor(square_2, cv2.COLOR_BGR2GRAY)
square_3 = cv2.cvtColor(square_3, cv2.COLOR_BGR2GRAY)
square_4 = cv2.cvtColor(square_4, cv2.COLOR_BGR2GRAY)
width, height = square_1.shape
square_1 = square_1 / 255.0
square_2 = square_2 / 255.0
square_3 = square_3 / 255.0
square_4 = square_4 / 255.0
pieces = ['knight', 'rook', 'bishop']
for pos, square in occupied:
def predict_board(occupied_squares: Squares) -> Board:
board = Board()
for pos, square in occupied_squares.items():
square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
width, height = square.shape
square = square / 255.0
test = new_model.predict(np.array(square).reshape((-1, width, height, 1)))
text_color = 255
cv2.putText(square, f"{pos} {pieces[int(np.argmax(test))]}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
color=(text_color,) * 3, thickness=3)
cv2.imwrite("lel", square)
cv2.waitKey(0)
"""
for pos, square in [(pos_1, square_1), (pos_2, square_2), (pos_3, square_3), (pos_4, square_4)]:
test = new_model.predict(np.array(square).reshape((-1, width, height, 1)))
print(f"{pos}: {np.argmax(test)}")
text_color = 255
cv2.putText(square, f"{pos} {pieces[int(np.argmax(test))]}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
color=(text_color,) * 3, thickness=3)
cv2.imshow(f"{pos}", square)
"""
#cv2.putText(square, f"{pos} {PIECE(int(np.argmax(test)))}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255,) * 3, thickness=3)
#cv2.imwrite("lel", square)
board[pos] = PIECE(int(np.argmax(test)))
return board

View File

@ -1,52 +1,35 @@
from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
import glob
import numpy as np
import cv2
#(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
#print(train_images[0])
import numpy as np
import tensorflow as tf
# exit()
from util import OUR_PIECES
# (train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
# print(train_images[0])
training_img = []
training_labels = []
test_img = []
test_labels_ = []
for piece in OUR_PIECES:
# training set
for _ in range(10):
for filename in glob.glob(f"../training_images/rook/*_square/*.png")[:-50]:
for filename in glob.glob(f"../training_images/{piece}/*_square/*.png")[:-50]:
training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
training_labels.append(1)
for _ in range(10):
for filename in glob.glob(f"../training_images/knight/*_square/*.png")[:-50]:
training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
training_labels.append(0)
for _ in range(10):
for filename in glob.glob(f"../training_images/bishop/*_square/*.png")[:-50]:
training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
training_labels.append(2)
training_labels.append(piece)
# test set
for _ in range(5):
for filename in glob.glob(f"../training_images/rook/*_square/*.png")[-50:]:
for filename in glob.glob(f"../training_images/{piece}/*_square/*.png")[-50:]:
test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
test_labels_.append(1)
for _ in range(5):
for filename in glob.glob(f"../training_images/bishop/*_square/*.png")[-50:]:
test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
test_labels_.append(2)
for _ in range(5):
for filename in glob.glob(f"../training_images/knight/*_square/*.png")[-50:]:
test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
test_labels_.append(0)
test_labels_.append(piece)
width, height = training_img[0].shape

26
util.py
View File

@ -3,7 +3,7 @@ from __future__ import annotations
from enum import Enum
from functools import lru_cache
from pathlib import Path
from typing import NamedTuple, Dict, Tuple
from typing import NamedTuple, Dict, Tuple, List
import cv2
import numpy as np
@ -19,23 +19,24 @@ class COLOR(Enum):
class PIECE(Enum):
PAWN = "pawn"
ROOK = "rook"
KNIGHT = "knight"
BISHOP = "bishop"
QUEEN = "queen"
KING = "king"
EMPTY = "empty"
KNIGHT = 0
ROOK = 1
BISHOP = 2
PAWN = 3
QUEEN = 4
KING = 5
EMPTY = 6
def __str__(self) -> str:
return self.value
return self.name.lower()
PieceAndColor = Tuple[PIECE, COLOR]
OUR_PIECES = (
PIECE.ROOK,
PIECE.KNIGHT,
PIECE.ROOK,
PIECE.BISHOP,
)
@ -85,6 +86,11 @@ Squares = Dict[POSITION, np.ndarray]
class Board(Dict[POSITION, PIECE]):
"""Board is a dict mapping positions to a piece, i.e. a board configuration after all image processing"""
@property
def to_array(self) -> List[List[int]]:
return [[self.get(POSITION((file, rank)), PIECE.EMPTY).value for file in FILE]
for rank in RANK]
def imwrite(*args, **kwargs):
Path(args[0]).parent.mkdir(parents=True, exist_ok=True)