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

# Load the model
model_path = "anime.pkl"  
learn = load_learner(model_path)

# Define categories
categories = learn.dls.vocab

# Define the image classification function
def Classify_image(img):
    pred,idx,probs = learn.predict(img)
    return dict(zip(categories, map(float, probs)))

# Define the Gradio interface with Blocks
with gr.Blocks() as demo:
    gr.Markdown("# Anime-character Classifier")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image()
            classify_button = gr.Button("Classify")
        with gr.Column():
            output_label = gr.Label(label="Classification Result")

    classify_button.click(Classify_image, inputs=image_input, outputs=output_label)

demo.launch()