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

model = from_pretrained_keras("merve/riddikulus")

labels = {0:"ravenclaw", 1:"gryffindor", 2:"slytherin",3:"ravenclaw",
          4:"ravenclaw", 5:"gryffindor", 6: "hufflepuff", 7:"ravenclaw",
          8:"ravenclaw", 9:"slytherin"}
          
canvas = gr.inputs.Image(source="canvas", shape=(28,28))          
text = gr.outputs.Textbox()          
def infer(image):
  cls = np.argmax(model.predict(np.expand_dims(image, axis = 0)[:,:,:,1]))
  cls = labels[cls]
  return f"Welcome to {cls}"
  
gr.Interface(infer, inputs=[canvas], outputs=[text]).launch(debug=True)