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

# Load model
learn = load_learner("model.pkl")

# Inference function
def classify_image(img):
    pred, pred_idx, probs = learn.predict(img)
    return {str(c): float(probs[i]) for i, c in enumerate(learn.dls.vocab)}

# Gradio interface
with gr.Blocks(title="Car body style classifier") as demo:
    gr.Markdown("# Upload a car image to classify its body style!")
    gr.Markdown("Uses `convnext_tiny` architecture and achieves *89.66% accuracy*.")
    gr.Markdown("This project was inspired by first two lectures of the [Practical Deep Learning for Coders](https://course.fast.ai/) course.")
    gr.Markdown("Trained [here](https://colab.research.google.com/drive/1wn4-22c1XopPIhM3uBW2Z6hAEAAHGozM)")

    with gr.Row():
        with gr.Column():
            inp = gr.Image(
                label="Upload a car image",
                type="pil"
            )
            btn = gr.Button("Submit")

        with gr.Column():
            out = gr.Label(num_top_classes=3)

    btn.click(classify_image, inputs=inp, outputs=out)

if __name__ == "__main__":
    demo.launch()