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

def predict_image(image):
    learner = learner_load('export.pkl')
    img = PILImage.create(image)
    pred = learner.predict(img)
    return pred[0]


def create_interface():
    
    image_input = gr.inputs.Image()

    output = gr.outputs.Text()

    iface = gr.Interface(
        fn=predict_image,
        inputs=image_input,
        outputs=output,
        capture_session=True,
        title="Cat or Dog Classifier",
        description="Upload an image to classify whether it's a cat or a dog."
    )
    return iface

if __name__ == "__main__":
    iface = create_interface()
    iface.launch()