advancedskrald/adapter.py

import base64
import json
import sys

import cv2
import numpy as np

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

while True:
    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)

    # 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))