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

model = tf.keras.models.load_model("hf://JaviSwift/cifar10_simple")

def predict_image(img):
    """
    Makes a prediction of the image descripton
    """

    img = tf.image.resize(img, (32, 32))
    img = img / 255.0
    img = np.expand_dims(img, axis=0)

    prediction = model.predict(img)

    predicted_class = np.argmax(prediction)

    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
    predicted_label = class_names[predicted_class]

    return predicted_label

iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(label="Upload an image"),
    outputs=gr.Label(label="Result"),
    title="Image description predictor",
    description="Upload an image and it will make a description of the object"
)

iface.launch()