import gradio as gr
import tensorflow as tf
import numpy as np

num_classes = 200
IMG_HEIGHT = 300
IMG_WIDTH = 300

with open("classlabel.txt", 'r') as file:
    CLASS_LABEL = [x.strip() for x in file.readlines()]

def normalize_image(img):
    img = tf.cast(img, tf.float32) / 255.0
    img = tf.image.resize(img, (IMG_HEIGHT, IMG_WIDTH), method='bilinear')
    return img

def load_model(model_name):
    # Load the model based on the model_name input
    if model_name == "model1":
        return tf.keras.models.load_model("model1.h5")
    elif model_name == "model2":
        return tf.keras.models.load_model("model2.h5")
    elif model_name == "model3":
        return tf.keras.models.load_model("model3.h5")
    else:
        raise ValueError("Invalid model_name")

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 {CLASS_LABEL[i]: str(v) for i, v in zip(top5_idx, top5_value)}

interface = gr.Interface(
    predict_top_classes,
    [
        gr.inputs.Image(type='pil'),
        gr.inputs.Button(label="Model 1 (Xception)", value="model1"),
        gr.inputs.Button(label="Model 2 (InceptionV3)", value="model2"),
        gr.inputs.Button(label="Model 3 (InceptionResNetV2)", value="model3"),
    ],
    outputs='label'
)
interface.launch()