import gradio as gr
from fastai.vision.all import *
import requests
from io import BytesIO


categories = ('Not Hot Dog and Not Pizza!', 'Hot Dog', 'Pizza')


def predict_category_of_image(url):
    response = requests.get(url)
    image = PILImage.create(BytesIO(response.content))
    learn = load_learner('model.pkl')
    predicted_category, _, probs = learn.predict(image)
    if predicted_category == "food":
        predicted_category = "not hot dog and not pizza"

    print(predicted_category)
    print(probs)
    print(dict(zip(categories, map(float, probs))))
    return [predicted_category, dict(zip(categories, map(float, probs))), image]


label = gr.outputs.Label(label="Confidence")
textbox = gr.outputs.Textbox(label="Predicted Category")
image = gr.outputs.Image(label="Image")

interface = gr.Interface(
    fn=predict_category_of_image,
    inputs="text",
    outputs=[textbox, label, image],
    title="Hot Dog or Not Hot Dog...or Pizza!"
)
interface.launch()

# why does it identify this as a hot dog? https://static01.nyt.com/images/2021/02/17/dining/17tootired-grilled-cheese/17tootired-grilled-cheese-articleLarge.jpg?quality=75&auto=webp&disable=upscale
# and this as a hot dog? https://tmbidigitalassetsazure.blob.core.windows.net/rms3-prod/attachments/37/1200x1200/Pizza-from-Scratch_EXPS_FT20_8621_F_0505_1_home.jpg
# example app: https://huggingface.co/spaces/jph00/testing