#!/usr/bin/env python3
''' Written by jared.vasquez@yale.edu '''

from keras.models import load_model
import numpy as np
import copy
import cv2
import os

dataColor = (0, 255, 0)
font = cv2.FONT_HERSHEY_SIMPLEX
fx, fy, fh = 10, 50, 45
showMask = 0
# Describe your own working directory.
# It can also be retrieved by "os.getcwd()"
DIR_DATA = r"D:\PROJECTS\GITHUB & KAGGLE\Github\CNN-HowManyFingers\images\train"

classes = 'NONE ONE TWO THREE FOUR FIVE'.split()


def initClass(name):
    os.system('mkdir -p data/%s' % name)


def binaryMask(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.GaussianBlur(img, (7, 7), 3)
    img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
    ret, new = cv2.threshold(img, 25, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    return new


def main():
    global font, fx, fy, fh
    global dataColor
    global showMask

    model = load_model('model_CNN-HowManyFingers.h5')

    x0, y0, width = 200, 220, 300

    cam = cv2.VideoCapture(0)
    cv2.namedWindow('Original', cv2.WINDOW_NORMAL)

    while True:
        # Get camera frame
        ret, frame = cam.read()
        frame = cv2.flip(frame, 1)  # mirror
        window = copy.deepcopy(frame)
        cv2.rectangle(window, (x0, y0), (x0 + width - 1, y0 + width - 1), dataColor, 12)

        # get region of interest
        roi = frame[y0:y0 + width, x0:x0 + width]
        roi = binaryMask(roi)

        # apply processed roi in frame
        if showMask:
            window[y0:y0 + width, x0:x0 + width] = cv2.cvtColor(roi, cv2.COLOR_GRAY2BGR)

        img = np.float32(roi) / 255.
        img = np.expand_dims(img, axis=0)
        img = np.expand_dims(img, axis=-1)
        pred = classes[int(np.argmax(model.predict([img])[0]))]
        cv2.putText(window, 'Prediction: %s' % pred, (fx, fy + 2 * fh), font, 1.0, (245, 210, 65), 2, 1)

        # show the window
        cv2.imshow('Original', window)

        # Keyboard inputs
        key = cv2.waitKey(10) & 0xff

        # use q key to close the program
        if key == ord('q'):
            break

        elif key == ord('b'):
            showMask = not showMask

    cam.release()


if __name__ == '__main__':
    initClass('NONE')
    main()