# Import libraries
import gradio as gr
import tensorflow as tf
import numpy as np

# Initialize the number of classes, also the image's height and width
num_classes = 200
IMG_HEIGHT = 300
IMG_WIDTH = 300

# Open the classlabel.txt to read the class labels
with open("classlabel.txt", 'r') as file:
    CLASS_LABEL = [x.strip() for x in file.readlines()]

# Function to normalize the image
def normalize_image(img):
    img = tf.cast(img, tf.float32)/255.
    img = tf.image.resize(img, (IMG_HEIGHT, IMG_WIDTH), method='bilinear')
    return img

# Function to select and load the model
def load_model(model_name):
    # Load the model based on the model_name input
    if model_name == "Xception":
        return tf.keras.models.load_model("model/Xception.h5")
    elif model_name == "InceptionV3":
        return tf.keras.models.load_model("model/InceptionV3.h5")
    elif model_name == "InceptionResNetV2":
        return tf.keras.models.load_model("model/InceptionResNetV2.h5")
    elif model_name == "DenseNet201":
        return tf.keras.models.load_model("model/DenseNet201.h5")
    else:
        raise ValueError("Invalid model_name")

# Main function, let the model make the prediction on the image uploaded
def predict_top_classes(img, model_name):
    img = img.convert('RGB')
    img_data = normalize_image(img)
    x = np.array(img_data)
    x = np.expand_dims(x, axis=0)
    model = load_model(model_name)
    temp = model.predict(x)

    idx = np.argsort(np.squeeze(temp))[::-1]
    top5_value = np.asarray([temp[0][i] for i in idx[0:5]])
    top5_idx = idx[0:5]
    
    # Return the top 5 highest probability class labels
    return {CLASS_LABEL[i]: str(v) for i, v in zip(top5_idx, top5_value)}

# Define the interface
interface = gr.Interface(
    predict_top_classes,
    [
        gr.Image(type='pil'),
        gr.Dropdown(
            choices=["Xception","InceptionV3","InceptionResNetV2","DenseNet201"],
            type="value",
            label="Select a model"
        )
    ],
    outputs='label'
)
    
# Launch the interface
interface.launch()