import gradio as gr
import os
import fastbook
from pathlib import Path
from fastai.vision.widgets import *

fastbook.setup_book()

examples = ['./examples/pikachu.webp', './examples/charizard.webp',
            './examples/mewtwo.jpg', './examples/rayquaza.jpeg', './examples/cinderace.webp']

path = Path()
path.ls(file_exts='.pkl')
learn_inf = fastbook.load_learner(path/'pokemon-detector.pkl')

labels = learn_inf.dls.vocab

def pokemon_classifier(image):
    image = fastbook.PILImage.create(image)
    pred, pred_id, probs = learn_inf.predict(image)
    output = {labels[i]: float(probs[i]) for i in range(len(labels))}
    # limit the output to the top 5 results
    output = dict(sorted(output.items(), key=lambda item: item[1], reverse=True)[:5])
    return output


iface = gr.Interface(fn=pokemon_classifier, inputs="image", outputs="label", examples=examples)
iface.launch()