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

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

# Define the prediction function
def predict(img):
    try:
        # Create the image object
        img = PILImage.create(img)
        # Get predictions from the model
        pred, pred_idx, probs = learn.predict(img)
        # Fetch the labels dynamically from the model's vocabulary
        labels = learn.dls.vocab
        # Ensure probabilities are floats
        return {labels[i]: float(probs[i]) for i in range(len(labels))}

    except Exception as e:
        # Log the exception and return it as an error message
        print(f"An error occurred: {e}")
        return {"error": str(e)}

# Define the Gradio interface
title = "Interior Design Classifier"
description = "Upload an image of an interior design and get a prediction of the design style."
examples = ['1.jpeg', '2.jpg', '3.jpg']

# Set up Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title=title,
    description=description,
    examples=examples
)

interface.launch(share=True, debug=False)