import gradio as gr
import tensorflow as tf
import cv2
import numpy as np
# from transformers import pipeline
from keras_preprocessing.image import img_to_array

# pipeline = pipeline(task="dogs-cat-classification", model="carlos-pino/dogs-cats")
# Get model
model = './saves/dogs-cats.h5'

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

IMAGE_SIZE = 100
def predict(frame):
    img_frame = cv2.imread(frame)
    # Parse to gray
    img = cv2.cvtColor(img_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)

    prediction_cat = 0
    prediction_dog = 0
    predictions = []

    # Classification
    if prediction <= 0.5:
        prediction_cat = (1 - prediction) / 0.5
        prediction_dog = 1 - prediction_cat
    else:
        prediction_cat = (prediction - 1) / 0.5
        prediction_dog = (prediction - 0.5) / 0.5

    predictions.append(dict(zip(["label", "score"], ["Cat", str(prediction_cat)])))
    predictions.append(dict(zip(["label", "score"], ["Dog", str(prediction_dog)])))
    return {p["label"]: p["score"] for p in predictions}


gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Upload cat or dog candidate", type="filepath"),
    outputs=gr.outputs.Label(num_top_classes=2),
    title="Dog or Cat?",
).launch()