# AUTOGENERATED! DO NOT EDIT! File to edit: app.ipynb.

# %% auto 0
__all__ = ['learn', 'categories', 'examples', 'intf', 'classify_image']

# %% app.ipynb 2
from fastai.vision.all import *
import gradio as gr

# Helpers used while building the model
# def is_cat(x): return x[0].isupper()

# %% app.ipynb 4
learn = load_learner("who_is_the_hero_model.pkl")

# %% app.ipynb 6
categories = learn.dls.vocab
categories = [category.capitalize() for category in categories]
print (f"Categories: {categories}")

def classify_image(img):
    pred, idx, probs = learn.predict(img)
    
    prediction = dict(zip(categories, map(float, probs)))
    print (f"prediction = {prediction}")
    
    predicted_hero = max(prediction, key=lambda key: prediction[key])
    print (f"predicted_hero = {predicted_hero}")
    
    if predicted_hero == 'Superman':
        alter_ego = "Clark Kent Jr"
    elif predicted_hero == "Batman":
        alter_ego = "Bruce Wayne"
    elif predicted_hero == "Flash":
        alter_ego = "Barry Allen"
    else:
        alter_ego = None
    
    return prediction, alter_ego

# %% app.ipynb 9
examples = [
    'images/batman.jpg', 'images/batman2.jpg', 'images/batman3.png', 
    'images/superman1.jpg', 'images/superman2.jpg', 'images/superman3.jpg',
    'images/flash1.jpg', 'images/flash2.jpg', 'images/flash3.jpg'
]

intf = gr.Interface(
    fn=classify_image, 
    inputs=gr.Image(shape=(192,192)), 
    outputs=[gr.Label(label='Predicted output'), gr.Text(label="Alter Ego")], 
    examples=examples,
    title="Who is the 'Super Hero' Classifier",
    description="Classifier is fine-tuned on pre-trained **resnet18** model using ~200 images in total"
)
intf.launch(inline=True)