Now we works and we still have type hints so thats nice
This commit is contained in:
parent
0f702640a3
commit
417538d768
|
@ -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
|
||||
|
|
14
adapter.py
14
adapter.py
|
@ -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))
|
||||
|
|
Binary file not shown.
Binary file not shown.
13
main.py
13
main.py
|
@ -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
|
||||
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
#(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
|
||||
# exit()
|
||||
from util import OUR_PIECES
|
||||
|
||||
#print(train_images[0])
|
||||
# (train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
|
||||
# print(train_images[0])
|
||||
|
||||
#exit()
|
||||
training_img = []
|
||||
training_labels = []
|
||||
|
||||
test_img = []
|
||||
test_labels_ = []
|
||||
|
||||
|
||||
for _ in range(10):
|
||||
for filename in glob.glob(f"../training_images/rook/*_square/*.png")[:-50]:
|
||||
for piece in OUR_PIECES:
|
||||
# training set
|
||||
for _ in range(10):
|
||||
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)
|
||||
training_labels.append(piece)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
for _ in range(5):
|
||||
for filename in glob.glob(f"../training_images/rook/*_square/*.png")[-50:]:
|
||||
# test set
|
||||
for _ in range(5):
|
||||
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
26
util.py
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue
Block a user