Update app.py

app.py

@@ -23,12 +23,16 @@ def predict(image):
     preprocessed = preprocess_image(image)
     # Make predictions
     predictions = model.predict(preprocessed)
-    # Get the predicted class index (with the highest probability)
-    predicted_class = np.argmax(predictions, axis=-1)[0]
+    
     # Define the class labels (you can update these with your actual class names)
     class_labels = ['Class 1', 'Class 2', 'Class 3', 'Class 4', 'Class 5']
-    # Return the predicted class with the highest probability
-    return class_labels[predicted_class], predictions[0][predicted_class]
+    
+    # Create a dictionary with class labels as keys and probabilities as values
+    prediction_dict = {class_labels[i]: float(predictions[0][i]) for i in range(len(class_labels))}
+    
+    # Get the predicted class with the highest probability
+    predicted_class = class_labels[np.argmax(predictions, axis=-1)[0]]
+    return predicted_class, prediction_dict
 
 # Gradio interface
 interface = gr.Interface(