advancedskrald/web.py

import base64

import cv2
import numpy as np
from flask import Flask, jsonify, request

from main import find_occupied_squares
from runner import find_homography, warp_board
from tensor_classifier import predict_board
from time import time

app = Flask(__name__)


@app.route("/", methods=["POST"])
def process():
    print("Received request")
    data = request.get_json(force=True)

    decoded = base64.b64decode(data["img"])
    img_array = np.frombuffer(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:
    start = time()
    print("Finding keypoints")
    homography = find_homography(camera_img, debug=True)
    print("Computing homography")
    warped_board = warp_board(camera_img, homography)
    print("Warping board")
    cv2.imwrite("warped.png", warped_board)
    print("Removing empty squares")
    occupied_squares = find_occupied_squares(warped_board)
    print("Predicting board state")
    board = predict_board(occupied_squares)
    print(f"The request took {round(time() - start, 3)} seconds")
    print("Returning board state")
    # Finally, output for unity to read
    return jsonify({
        "homography": homography.tolist(),
        "board": board.to_array,
        })


def main():
    app.run(host='0.0.0.0', debug=True)


if __name__ == '__main__':
    main()