import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.models import load_model

HEIGHT, WIDTH = 224, 224
IMG_SIZE = 224
model = load_model('Models/best_model1.h5')

def classify_image(inp):
    labels = ['Cat', 'Dog']
    inp = tf.image.resize(inp, [IMG_SIZE, IMG_SIZE])
    inp = inp.numpy().reshape((-1, IMG_SIZE, IMG_SIZE, 3))
    inp = tf.keras.applications.vgg16.preprocess_input(inp)
    prediction = model.predict(inp).flatten()
    if len(prediction) == 1:
        dog_prob = float(prediction[0])
        return {labels[0]: 1 - dog_prob, labels[1]: dog_prob}
    else:
        return {labels[i]: float(prediction[i]) for i in range(len(labels))}

image = gr.Image(height=HEIGHT, width=WIDTH, label='Input')
label = gr.Label(num_top_classes=2)
examples = [["Examples/img1.png"], ["Examples/img2.png"],["Examples/img3.png"], ["Examples/img4.png"]]


gr.Interface(
    fn=classify_image, 
    inputs=image, 
    outputs=label, 
    title='Smart Pet Classifier',
    examples=examples
).launch(debug=False)