import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np


def load_models():
    models = {}
    models['SimpleNN_model'] = tf.keras.models.load_model("Model_catsVSdogs.h5")
    models['VGG16'] = tf.keras.models.load_model("vgg16.h5")  
    return models


models = load_models()


def predict_image(img, model_name):
    model = models[model_name]
    
    if model_name == 'SimpleNN_model':
        img = img.resize((256, 256))
    elif model_name == 'VGG16':
        img = img.resize((224, 224))

    img_array = image.img_to_array(img)  
    img_array = np.expand_dims(img_array, axis=0)  
    img_array = img_array / 255.0  


    prediction = model.predict(img_array)
    if prediction[0] < 0.5:
        return "Cat"
    else:
        return "Dog"


interface = gr.Interface(fn=predict_image, 
                         inputs=[gr.Image(type="pil"), gr.Dropdown(["SimpleNN_model", "VGG16"], label="Select Model")], 
                         outputs="text",
                         title="Cat and Dog Classifier",
                         description="Upload an Image")


interface.launch(share=True)