import warnings
warnings.filterwarnings("ignore")

import pathlib
pathlib.WindowsPath = pathlib.PosixPath

import gradio as gr
from fastbook import *
from fastai.vision.widgets import *


model = load_learner('./hero_model.pkl')

categories = ('Batman', 'Hulk', 'Ironman', 'Spiderman', 'Superman', 'Wonderwoman')

def classify_img(img):
    pred, idx, probs = model.predict(img)
    return dict(zip(categories, map(float, probs)))

image = gr.Image(width=512, height=512)
label = gr.Label()
examples = ['batman.jpg', 'superman.jpg', 'ironman.jpg', 'spiderman.jpg', 'hulk.jpg', 'wonderwoman.jpg']

gr.Interface(fn=classify_img, inputs=image, outputs=label, examples=examples, title="Hero Classifier",
    description="Is it batman? or superman? or ironman? or spiderman? or hulk? or wonderwoman?",).launch(inline=False)