Update app.py

app.py

@@ -4,28 +4,28 @@ import gradio as gr
 # Load the model
 learn = load_learner('model.pkl')
 
-# Define the prediction function
+# Fetch the labels dynamically from the model's vocabulary
 labels = learn.dls.vocab
 
-
+# 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)
-        # Convert boolean keys to strings and ensure probabilities are floats
-        return {"Bauhaus": float(probs[0]), "Bohemian": float(probs[1]), "Scandinavian": float(probs[2])}
+        # Convert probabilities to percentages and construct the output dictionary dynamically
+        return {labels[i]: f"{probs[i]*100:.2f}%" 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,