import gradio as gr
from fastbook import PILImage, load_learner

labels = [
    'Adam Sandler', 'Angelina Jolie', 'Brad Pitt', 'Chris Hemsworth',
    'Denzel Washington', 'Dwayne Johnson', 'George Clooney', 'Harrison Ford',
    'Johnny Depp', 'Julia Roberts', 'Keanu Reeves', 'Kevin Hart',
    'Leonardo DiCaprio', 'Morgan Freeman', 'Robert Downey Jr', 'Ryan Reynolds',
    'Sandra Bullock', 'Tom Cruise', 'Tom Hanks', 'Will Smith',
]


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


gr.Interface(
    predict,
    inputs=gr.inputs.Image(shape=(224, 224)),
    outputs=gr.outputs.Label(num_top_classes=20),
    title="Hollywood actor",
).launch()