Update app.py

app.py

@@ -4,6 +4,21 @@ import numpy as np
 from PIL import Image
 import urllib.request
 from utils import *  # Assuming the gen_labels() and preprocess() functions are in this module
+from your_model_module import model_arc  # Import your model function
+
+# Function to classify the garbage
+def classify_garbage(img_path, model):
+    processed_img = preprocess_image(img_path)
+    prediction = model.predict(processed_img)
+
+    class_labels = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
+    predicted_class = np.argmax(prediction, axis=1)[0]
+    classification_result = class_labels[predicted_class]
+
+    # Get the confidence (probability) of the predicted class
+    confidence = prediction[0][predicted_class] * 100  # Convert probability to percentage
+
+    return classification_result, confidence
 
 # Load labels
 labels = gen_labels()
@@ -31,15 +46,10 @@ if 'image' in locals():  # Check if image variable exists
 
     if st.button('Predict'):
         try:
-            img_array = preprocess(image)
-            model = model_arc()  # Assuming this function initializes and returns your model
-            prediction = model.predict(img_array[np.newaxis, ...])
-            
-            # Get the predicted class name
-            predicted_class_index = np.argmax(prediction[0], axis=-1)
-            predicted_class_name = labels[predicted_class_index]
+            model = model_arc()  # Initialize your model
+            predicted_class, confidence = classify_garbage(image, model)
             
-            st.info('The uploaded image has been classified as "{}" waste.'.format(predicted_class_name))
+            st.info('The uploaded image has been classified as "{}" waste with {:.2f}% confidence.'.format(predicted_class, confidence))
             
         except Exception as e:
             st.error(f"An error occurred: {e}")