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

path = Path()
path.ls(file_exts='.pkl')

# Specify the path to your fastai model file
learn_inf = load_learner(path/'export.pkl')

# Define the predict function
def predict(image):
    # Load the input image
    img = PILImage.create(image)

    # Perform inference using the loaded fastai Learner
    pred, pred_idx, probs = learn_inf.predict(img)

    return f'Prediction: {pred}; Probability: {probs[pred_idx]:.04f}'

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(upload_button=True, label="Upload a picture of a dog"),
    outputs="text",
    live=True,
    title="Dog Breed Classifier",
    description="Upload a picture of a dog, and the model will classify it as Labrador, Husky, or Bulldog.",
    examples=["Axel.jpg"])


# Launch the Gradio interface
iface.launch()