import gradio as gr
import numpy as np
from numpy import asarray
from PIL import Image
from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras('keras-io/SimSiam')

index_to_name = {0:'Airplane', 1:'Car', 2:'Bird',
                 3:'Cat', 4:'Deer', 5:'Dog',
                 6:'Frog', 7:'Horse', 8:'Ship',
                 9:'Truck'}
                 
def predict_with_simsiam(original_image): 
    image = asarray(original_image)
    image = np.expand_dims(image, axis=0)
    pred_prob = model.predict(image).flatten().tolist()
    return {index_to_name[i]: pred_prob[i] for i in range(10)}
    
                     
title = "Self-supervised contrastive learning with SimSiam"
description = "This space implements a SimSiam network for image classification of the Cifar 10 dataset."
examples = ['horse1.png', 'airplane4.png', 'dog6.png']
article = """<p style='text-align: center'>
        <a href='https://keras.io/examples/vision/simsiam' target='_blank'>Keras Example given by Sayak Paul</a>
        <br>
        Space by @Jezia
    </p>
    """

iface = gr.Interface(predict_with_simsiam, inputs=[gr.inputs.Image(label="image", type="pil")], outputs="label", title=title, description=description, article=article, examples=examples)
iface.launch(debug='True')