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

import fastai.losses
import fastai.layers

fastai.layers.BaseLoss = fastai.losses.BaseLoss
fastai.layers.CrossEntropyLossFlat = fastai.losses.CrossEntropyLossFlat
fastai.layers.BCEWithLogitsLossFlat = fastai.losses.BCEWithLogitsLossFlat
fastai.layers.BCELossFlat = fastai.losses.BCELossFlat
fastai.layers.MSELossFlat = fastai.losses.MSELossFlat
fastai.layers.L1LossFlat = fastai.losses.L1LossFlat
fastai.layers.LabelSmoothingCrossEntropy = fastai.losses.LabelSmoothingCrossEntropy
fastai.layers.LabelSmoothingCrossEntropyFlat = fastai.losses.LabelSmoothingCrossEntropyFlat

model = load_learner("model.pkl")
def predict(im):
    image_file = PILImage(PILImage.create((255-im)))
    pred,pred_idx,probs = model.predict(image_file)
    vals, indx = torch.topk(probs,2)
    
    return {model.dls.vocab[i]: prob.item() for prob,i in zip(vals,indx)}

input_widget = gr.inputs.Image(image_mode="L", source="canvas", shape=((224,224)), invert_colors=True)
# 
# <b><span style="color:#fab114">some *blue* text</span>.</b>
classes = ",".join(model.dls.vocab)
article = f'currently supports <b><span style="color:#fab114">{classes}</span>.</b>'
interface = gr.Interface(predict, title="Quickdraw", inputs=input_widget, outputs='label', live=True,article=article)

interface.launch(debug=True)