import gradio as gr
from keras.models import load_model
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
from tensorflow.keras.utils import img_to_array
from tensorflow.keras.applications.resnet50 import preprocess_input
from PIL import Image

# model path
cnn_model = load_model('./model/cnn_model.h5')
resnet_model = load_model('./model/resnet_model.h5')

import json

with open('data/class_dict.json', 'r') as json_file:
    class_dict = json.load(json_file)

# get class names
class_names = [class_dict[i] for i in sorted(class_dict.keys())]


def classify_insect(model_name, img):
    img = img.resize((150, 150))
    img_array = img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_preprocessed = preprocess_input(img_array)

    # Select the model based on the dropdown choice
    if model_name == "CNN Model":
        model = cnn_model
    elif model_name == "Transfer Learning ResNet":
        model = resnet_model

    # Make a prediction
    prediction = model.predict(img_preprocessed)
    return {class_name: float(score) for class_name, score in zip(class_names, prediction[0])}

iface = gr.Interface(
    fn=classify_insect,
    inputs=[
        gr.Dropdown(choices=["CNN Model", "Transfer Learning ResNet"], label="Select Model"),
        gr.Image(shape=(150,150))
    ],
    outputs=gr.Label(num_top_classes=3)
)

iface.launch(share=True)