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

# Load the model
learn = load_learner('model.pkl')

# Define prediction function
def classify_image(input_img):
    # Convert input image to FastAI-compatible format
    input_img = PILImage.create(input_img)
    pred, idx, probs = learn.predict(input_img)
    return input_img, {learn.dls.vocab[i]: float(probs[i]) for i in range(len(probs))}

# Define Gradio app
gradio_app = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(label="Upload Image", sources=['upload', 'webcam'], type="pil"),
    outputs=[
        gr.Image(label="Processed Image"),
        gr.Label(label="Prediction Results", num_top_classes=5)
    ],
    title="Image Classification App",
    examples=["basset.jpg"]
)

# Launch the app
if __name__ == "__main__":
    gradio_app.launch()