Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import numpy as np
+import matplotlib.pyplot as plt
 from PIL import Image
 import tensorflow as tf
 from tensorflow.keras.preprocessing import image
@@ -31,13 +32,29 @@ def classify_image(pil_image):
     
     # Get predicted class and confidence
     predicted_class = np.argmax(prediction)
-    predicted_class = waste_labels[predicted_class]
+    predicted_class_name = waste_labels[predicted_class]
     confidence = prediction[0, np.argmax(prediction)]
 
-    return predicted_class
+    # Get class names and probabilities for all classes
+    class_names = list(waste_labels.values())
+    probabilities = prediction[0]
+
+    # Plot bar chart and save it to a file
+    plt.bar(class_names, probabilities, color='blue')
+    plt.xlabel('Waste Classes')
+    plt.ylabel('Probability')
+    plt.title('Prediction Probabilities')
+    plt.savefig('prediction_plot.png')
+
+    # Prepare output text with all class probabilities
+    output_text = f"Predicted Class: {predicted_class_name}, Confidence: {confidence:.4f}\n"
+    for class_name, prob in zip(class_names, probabilities):
+        output_text += f"{class_name}: {prob:.4f}\n"
+
+    return output_text, 'prediction_plot.png'
 
 # Create the Gradio interface
-iface = gr.Interface(fn=classify_image, inputs="image", outputs="text")
+iface = gr.Interface(fn=classify_image, inputs="image", outputs=["text", "image"],live=True)
 
 # Launch the Gradio interface
 iface.launch()