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:"Ravenclaw 🦅💙",3:"Slytherin🐍💚",
          4:"Hufflepuff 🦡💛", 5:"Death eater detected! ", 6: "Hufflepuff 🦡💛", 7:"Slytherin🐍💚",
          8:"Ravenclaw 🦅💙", 9:"Gryffindor 🦁"}
          
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]))
  if cls == 5:
    output = "Death eater detected! 💀"
  else:
    cls = labels[cls]
    output = f"Welcome to {cls}"
  return output
  
gr.Interface(infer, inputs=[canvas], outputs=[text], title="Welcome to Hogwarts Sorting Hat!", description="Draw something and let the sorting hat sort you! 🎩 ").launch()