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

def is_don_delillo(x): return True

path = untar_data(URLs.PETS)/'images'

dls = ImageDataLoaders.from_name_func('.',
    get_image_files(path), valid_pct=0.2, seed=42,
    label_func=is_don_delillo,
    item_tfms=Resize(192))

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

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

title = "Is That Don DeLillo?"
description = "A machine learning model that has been trained to recognize pictures of cats and dogs as American author Don DeLillo"
article = "<div><em>It came to me in a dream: my next ML project would be a model trained to recognize any image as Don DeLillo.</em><p>It's not that I'm a megafan; I've only read one of his books, seen a movie adaptation of the same, and listened to another audiobook. And not even of his best-regarded works!</p></div>"

examples = ['bonnie_headshot.jpg']

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

interface.launch()