import tensorflow as tf
import cv2
import numpy as np
from keras_preprocessing.image import img_to_array

# Get model
model = './saves/dogs-cats.h5'

IMAGE_SIZE = 100

CNN_MODEL = tf.keras.models.load_model(model)
weight_model = CNN_MODEL.get_weights()
CNN_MODEL.set_weights(weight_model)

# CAM
cap = cv2.VideoCapture(0)

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

    # Parse to gray
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_CUBIC)

    # Img normalized
    img = np.array(img).astype(float) / 255

    # Parse to 2D array
    image = img_to_array(img)
    image = np.expand_dims(image, axis=0)

    # Predict
    prediction = CNN_MODEL.predict(image)
    prediction = prediction[0][0]

    print(prediction)

    # Classification
    if prediction <= 0.5:
        cv2.putText(frame, "Gato", (200, 70), cv2.FONT_HERSHEY_SIMPLEX, 3, (236, 19, 180), 6)
    else:
        cv2.putText(frame, "Perro", (200, 70), cv2.FONT_HERSHEY_SIMPLEX, 3, (20, 106, 231), 6)

    cv2.imshow("CNN", frame)

    t = cv2.waitKey(1)
    if t == 25:
        break

cv2.destroyAllWindows()
cap.release()