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

categories = ("grizzly bear", "black bear", "teddy bear")
learn = load_learner("export.pkl")

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

with gr.Blocks() as demo:
    gr.Markdown("# 🐻 Bear Classifier")
    gr.Markdown("Classify different types of bears: grizzly, black, or teddy bears.")

    with gr.Row():
        image_input = gr.Image(height=192, width=192)
        label_output = gr.Label()

    btn = gr.Button("Classify")
    btn.click(fn=classify_image, inputs=image_input, outputs=label_output)

    gr.Examples(
        examples=[
            ["examples/bears.jpg"],
            ["examples/grizzly.jpg"]
        ],
        inputs=image_input
    )

if __name__ == "__main__":
    print(f"===》Torch version: {torch.__version__}")
    print(f"===》Fastai version: {fastai.__version__}")
    print(f"===》Fastcore version: {fastcore.__version__}")
    print(f"===》Python version: {sys.version}")
    demo.launch()