import gradio as gr
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
import numpy as np

model = from_pretrained_keras("keras-io/semi-supervised-classification-simclr")

labels_gradio = ["Avión", "Pajaro", "Coche", "Gato", "Ciervo", "Perro", "Caballo", "Mono", "Barco", "Camión"]


def predict(imput_image):
    image = tf.constant(imput_image)
    image = tf.reshape(image, [-1, 96, 96, 3])
    pred = model.predict(image)
    pred_list = pred[0, :]
    pred_softmax = np.exp(pred_list)/np.sum(np.exp(pred_list))
    softmax_list = pred_softmax.tolist()
    return {labels_gradio[i]: softmax_list[i] for i in range(10)}


image = gr.inputs.Image(shape=(96, 96))
label = gr.outputs.Label(num_top_classes=4)

pie_pag = """<center>
Modelo: <a href='https://huggingface.co/keras-io/semi-supervised-classification-simclr' target='_blank'>keras.io</a>
Basado en el Space: <a href='https://huggingface.co/spaces/keras-io/semi-supervised-classification' target='_blank'>András Béres</a>
Autor: <a href='https://huggingface.co/machde' target='_blank'>Manuel Chacón De Dios</a>"""

titulo = "Mini clasificador"
descripcion = """<center>Clasificador capaz de etiquetar si es un Avión, Pajaro, Coche, 
Gato, Ciervo, Perro, Caballo, Mono, Barco, Camión</center>"""




Iface = gr.Interface(
    fn=predict,
    inputs=image,
    outputs=label,
    layout="vertical",
    theme="seafoam",
    examples=[["test_img/pajaro-test.jpeg"], ["test_img/coche-test.jpg"], ["test_img/perro-test.jpg"]],
    title=titulo,
    article=pie_pag,
    description=descripcion,
).launch()