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

model = tf.keras.models.load_model('JaviSwift/cifar10_simple.keras')

def predict_image(img):
    """
    Realiza la predicción sobre la imagen dada usando el modelo Keras.
    """

    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(shape=(32, 32), image_mode="RGB", source="upload", tool="editor"),
    outputs=gr.Label(num_top_classes=3),
    title="Clasificador de Imágenes con Keras",
    description="Cargue una imagen para clasificarla usando un modelo Keras entrenado en CIFAR-10." 
)

iface.launch()