StanKarz's picture
Update app.py
628f938 verified
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)