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


CATS_MAP = {
    "picasso": "Pablo Picasso",
    "vanGogh": "Vincent van Gogh",
    "dali": "Salvador Dalí",
    "daVinci": "Leonardo da Vinci",
    "rembrandt": "Rembrandt",
}

CATS_MAP_V2 = {
    "picasso": "Pablo Picasso",
    "vanGogh": "Vincent van Gogh",
    "dali": "Salvador Dalí",
    "daVinci": "Leonard da Vinci",
    "rembrandt": "Rembrandt",
    "monet": "Claude Monet",
    "caruso": "Santiago Caruso",
    "renoir": "Pierre-Auguste Renoir",
    "oKeeffe": "Georgia O’Keeffe",
    "krasner": "Lee Krasner",
}

# load pre-trained model
model = load_learner("model_v2.pkl")

# get classes name in right order
full_name_cats = [CATS_MAP_V2[key_class] for key_class in model.dls.vocab]


def classify_image(img) -> dict:
    category, idx, probs = model.predict(img)

    return dict(zip(full_name_cats, map(float, probs)))


# Gradio control
image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label(num_top_classes=4)
examples = [
    f"images_examples/{filename}" 
    for filename in ("mona_lisa.jpg", "starry_night.jpg", "persistence_memory.jpg")
]

gui = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    examples=examples,
    title="Detect the painter",
    description=(
        f"Detect if the given painting image is by a famous painter ({full_name_cats})."
    )
)
gui.launch(inline=False)