from fastai.vision.all import *
from huggingface_hub import from_pretrained_fastai
import gradio as gr

learn = from_pretrained_fastai("jeffreymjohnson/Deadpool-Detector")

labels = learn.dls.vocab

def predict(img):
    yes_reaction = "reactions/yes.jfif"
    no_reaction = "reactions/no.jfif"
    what_reaction = "reactions/what.jpg"
    
    img = PILImage.create(img)
    pred,pred_idx,probs = learn.predict(img)
    index = pred_idx.item()
    probability = probs[index].item()
    isDeadpool = pred == "Deadpool"
    
    if isDeadpool:
        x = "is"
    else:
        x = "is not"
    
    
    if probability < .75:
        return PILImage.create(what_reaction)
    elif isDeadpool:
        return PILImage.create(yes_reaction)
    else:
        return PILImage.create(no_reaction)
    
demo = gr.Interface(
    title="Deadpool™️ Detector",
    fn=predict,
    inputs=gr.Image(shape=(200,200)),
    outputs=gr.Image(shape=(200,200)), 
    examples=["examples/deadpool_example1.jpg", "examples/deadpool_example2.jpg", "examples/deadpool_example3.jpg", "examples/spiderman.jpg",]
).launch()