#from PIL import Image
from sys import argv, exit
import numpy as np
import cv2 as cv

from tensorflow.keras.utils import array_to_img as img

from matplotlib import pyplot as plt

if len(argv) != 2:
    exit(1)

#i = np.array(Image.open(argv[1]).convert('RGB').resize((300,300)))

from test2 import *
m.load_weights("model3.keras")

c = cv.VideoCapture(argv[1])

while True:
    ret, frame = c.read()

    if not ret:
        break

    pi = np.array(img(cv.cvtColor(frame, cv.COLOR_BGR2RGB)).resize((300, 300)).convert("RGB"))
    #plt.imshow(pi)
    #plt.show()

    #print(pi.astype(float))

    r = m.predict(np.array([pi.astype(float)]))

    if r[0][0] >= 0.5:
        if (pi[30][30][0] >= 190):
            cv.putText(frame, f"Y: {r[0][0]:.9f}", (10, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (0,40,0), 2)
        else:
            cv.putText(frame, f"Y: {r[0][0]:.9f}", (10, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (200,255,200), 2)
    else:
        if (pi[30][30][0] >= 190):
            cv.putText(frame, f"N: {r[0][0]:.9f}", (10, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (0,40,0), 2)
        else:
            cv.putText(frame, f"N: {r[0][0]:.9f}", (10, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (200,255,200), 2)

    cv.imshow('me', frame)

    if cv.waitKey(1) == ord('q'):
        break

#print(frame)

c.release()
cv.destroyAllWindows()

#r = mod.predict(np.array([i]))
#print(r)