import gradio as gr
from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras("AiresPucrs/digit-classifier")

classes = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine']

def predict(img):
  img = img.reshape(1, 28, 28, 1)
  prediction = model.predict(img, verbose=0).tolist()[0]
  return {classes[i]: prediction[i] for i in range(10)}

title = "Digit Classifier - By Teeny-Tiny Castle 🏰"

head = (
  "<center>"
  "<img src='https://upload.wikimedia.org/wikipedia/commons/2/27/MnistExamples.png' width=400>"
  "This model was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
  "</center>"
)


ref = (
"<center>"
"Return to the <a href='https://github.com/Nkluge-correa/teeny-tiny_castle)'>castle</a>."
"</center>")

# create interface
demo = gr.Interface(fn=predict,
             inputs="sketchpad",
             outputs=gr.Label(num_top_classes=3),
             allow_flagging="never",
             title=title,
             description=head,
             article=ref)

# launch interface
demo.launch()