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


learner = load_learner('./matt-damon-classifier.pkl')
label_map = {
    'neal-mcDonough': 'Neal McDonough',
    'mark-wahlberg': 'Mark Wahlberg',
    'jesse-plemons': 'Jesse Plemons',
    'matt-damon': 'Matt Damon'
}

def image_classifier(img):
    pred, _, probs = learner.predict(img)

    return label_map[pred]


iface = gr.Interface(fn=image_classifier, inputs="image", outputs="label", examples=[
    ["jesse-plemons.jpg"],
    ["jesse-plemons-2.webp"],
    ["mark-wahlberg.jpg"],
    ["matt-damon.jpg"],
    ["neal-mcDonough.jpg"]
    ],)
iface.launch()