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

learn = load_learner("model_2.pkl")

categories = learn.dls.vocab
for index, category in enumerate(categories):
    if category == "Random Anime Photos":
        categories[index] = "Others"


def classify_image(img):
    pred, idx, probs = learn.predict(img)
    return dict(zip(categories, map(float, probs)))


image = gr.Image()
label = gr.Label()
examples = [
    "Luffy.jpg",
    "Naruto-Kurama-Mode.png",
    "Goku.jpg",
    "Ichigo.jpeg",
    "Robin.jpeg",
]

title = "Top 10 Shounen Anime Protagonists Classifier"
description = "Fine tuned a resnet152 image classifier such that it is able to recognize protagonists of top 10 Shounen Animes."
start_article = (
    "<p> Animes and its protagonists this image classifier will recognize:</p>"
)
anime_characters = [
    "1. One Piece - Monkey D. Luffy<br>",
    "2. Naruto: Shippuden - Naruto Uzumaki<br>",
    "3. My Hero Academia - Izuku Midoriya<br>",
    "4. Dragon Ball Z - Son Goku aka Kakarot<br>",
    "5. Fairy Tail - Natsu Dragneel<br>",
    "6. Yu Yu Hakusho - Yusuke Urameshi<br>",
    "7. Bleach - Ichigo Kurosaki<br>",
    "8. Hunter X Hunter - Gon Freecss<br>",
    "9. Fullmetal Alchemist - Edward Elric<br>",
    "10. Attack on Titan - Eren Yeager<br>",
]
end_article = "<p>Rest all other anime characters will be classified as others.</p>"
final_article = start_article + "".join(anime_characters) + end_article

intf = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    examples=examples,
    title=title,
    description=description,
    article=final_article,
)
intf.launch(inline=False)