import numpy as np
import gradio as gr
import tensorflow as tf
from io import StringIO
from PIL import Image

labels = []
model = tf.keras.models.load_model('./models.h5')
with open("labels.txt") as f:
    for line in f:
        labels.append(line.replace('\n', ''))

def classify_image(inp):
    # Create a copy of the input array to avoid reference issues
    inp_copy = np.copy(inp)
    # Resize the input image to the expected shape (224, 224)
    inp_copy = Image.fromarray(inp_copy)
    inp_copy = inp_copy.resize((224, 224))
    inp_copy = np.array(inp_copy)
    inp_copy = inp_copy.reshape((-1, 224, 224, 3))
    inp_copy = tf.keras.applications.efficientnet.preprocess_input(inp_copy)
    prediction = model.predict(inp_copy).flatten()
    confidences = {labels[i]: float(prediction[i]) for i in range(90)}
    return confidences

demo = gr.Interface(classify_image, gr.Image(), gr.Label(num_top_classes=3))
if __name__ == "__main__":
    demo.launch()