advancedskrald/adapter.py

32 lines
809 B
Python
Raw Permalink Normal View History

import base64
2019-04-11 11:39:19 +00:00
import json
2019-04-16 21:29:41 +00:00
import sys
2019-04-11 11:39:19 +00:00
import cv2
import numpy as np
from main import find_occupied_squares
from runner import find_homography, warp_board
from tensor_classifier import predict_board
2019-04-24 07:41:57 +00:00
2019-04-25 14:00:04 +00:00
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)
2019-04-24 07:37:37 +00:00
2019-04-25 14:00:04 +00:00
# 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)
2019-04-24 07:37:37 +00:00
2019-04-25 14:00:04 +00:00
# Finally, output to stdout for unity to read
result = {
"homography": homography.tolist(),
"board": board.to_array,
}
2019-04-24 07:37:37 +00:00
2019-04-25 14:00:04 +00:00
print(json.dumps(result))