from fastai.vision.all import *
import gradio as gr

learner_madness = load_learner('madness.pkl')
categories = ['The Auditor', 'Deimos', 'Hank', "Jeb", 'Sanford', 'Tricky']

def classify_madness(img):
  prediction, index, probability = learner_madness.predict(img)
  return dict(zip(categories, map(float, probability)))

intf = gr.Interface(fn=classify_madness,
             inputs=gr.Image(shape=(192, 192)),
             outputs=gr.Label(),
             title="Which Madness Combat Character",
             description="This app was created for my 14 year old son. It can differentiate between The Auditor, Tricky, Hank, Sanford, Jeb, and Deimos! Try it out! Maybe see which it thinks you most resemble!",
             examples=['auditor.jpg', 'hank.jpg', 'tricky.jpg', 'jeb.jpg', 'sanford.jpg', 'deimos.jpg'])
intf.launch(inline=False)