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 = ["airplane", "bird", "car", "cat", "deer", "dog", "horse", "monkey", "ship", "truck"]


def infer(test_image):
    image = tf.constant(test_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[i]: softmax_list[i] for i in range(10)}


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


article = """<center>
Authors: <a href='https://twitter.com/johko990' target='_blank'>Johannes Kolbe</a> after an example by András Béres at 
<a href='https://keras.io/examples/vision/semisupervised_simclr/' target='_blank'>keras.io</a>"""


description = """Image classification with a model trained via Semi-supervised Contrastive Learning """


Iface = gr.Interface(
    fn=infer,
    inputs=image,
    outputs=label,
    examples=[["examples/monkey.jpeg"], ["examples/titanic.jpg"], ["examples/truck.jpg"]],
    title="Semi-Supervised Contrastive Learning Classification",
    article=article,
    description=description,
).launch()