Update app.py

app.py

@@ -10,18 +10,15 @@ def classify_garbage(img_path, model):
     processed_img = preprocess(img_path)
     prediction = model.predict(processed_img)
 
-    class_labels = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
+    labels = gen_labels()
     predicted_class = np.argmax(prediction, axis=1)[0]
-    classification_result = class_labels[predicted_class]
+    classification_result = 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()
-
 # Streamlit app layout
 st.markdown('<center><h1>Garbage Segregation</h1></center>', unsafe_allow_html=True)
 st.markdown('<center><h3>Please upload Waste Image to find its Category</h3></center>', unsafe_allow_html=True)
@@ -46,9 +43,13 @@ if 'image' in locals():  # Check if image variable exists
     if st.button('Predict'):
         try:
             model = model_arc()  # Initialize your model
-            predicted_class, confidence = classify_garbage(image, model)
             
-            st.info('The uploaded image has been classified as "{}" waste with {:.2f}% confidence.'.format(predicted_class, confidence))
+            # Ensure image shape is correct and add batch dimension
+            img_array = preprocess(image)  # This should return an array of shape (1, 256, 256, 3)
+            
+            predicted_class, confidence = classify_garbage(img_array, model)
+            
+            st.info('The uploaded image has been classified as "{}" with {:.2f}% confidence.'.format(predicted_class, confidence))
             
         except Exception as e:
             st.error(f"An error occurred: {e}")