from fastai.vision.all import *

learn = load_learner('poop.pkl')
labels = learn.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

import gradio as gr


# add a title and description
title = "Is this a photo of a doing his business? Or just a dog hanging out?"
desc = "My dog Kona has a nasty habit of doing her business in a very particular spot in the upstairs of our house. Using this model, I'll be able to aim a camera at that spot and send a text alert whenever she does this, so I'll have a better chance of correcting the bad behavior in the moment, which is essential to training a dog."


article = "Check out [my blog post](https://currentlyobsessed.com/p/my-first-deep-learning-model-doo-doo-detective) about this project."

examples = ['kona.jpg', 'baddog.jpg', 'kona2.jpg', 'kona3.jpg', 'baddog2.jpg']

enable_queue=True

iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(512, 512)), outputs=gr.outputs.Label(num_top_classes=3), title=title, description=desc, examples=examples, article=article)

iface.launch()