Update app.py

app.py

@@ -101,6 +101,7 @@ def train_model():
         MaxPooling2D(2,2),
 
         Flatten(),
+
         Dense(256, activation='relu'),
         Dropout(0.5),
 
@@ -133,8 +134,7 @@ if not os.path.exists(MODEL_PATH):
 else:
     model = tf.keras.models.load_model(MODEL_PATH)
 
-    # IMPORTANT:
-    # Ensure this matches dataset folder order exactly
+    # Ensure exact dataset folder order
     classes = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
 # -----------------------------
@@ -151,6 +151,9 @@ uploaded_file = st.file_uploader(
 
 if uploaded_file is not None:
     try:
+        # -----------------------------
+        # LOAD IMAGE
+        # -----------------------------
         image = Image.open(uploaded_file).convert("RGB")
 
         st.image(
@@ -170,44 +173,67 @@ if uploaded_file is not None:
         # PREDICT
         # -----------------------------
         with st.spinner("Analyzing waste type..."):
-            prediction = model.predict(img_array, verbose=0)
-
-        probabilities = prediction[0]
+            prediction = model.predict(
+                img_array,
+                verbose=0
+            )
 
-        predicted_index = np.argmax(probabilities)
-        predicted_class = classes[predicted_index]
-        confidence = probabilities[predicted_index] * 100
+        probabilities = prediction.flatten()
 
         # -----------------------------
-        # DISPLAY RESULTS
+        # VALIDATE OUTPUT
         # -----------------------------
-        st.subheader("📊 Prediction Scores")
+        if len(probabilities) != len(classes):
+            st.error(
+                f"Model output mismatch: Expected {len(classes)} classes but got {len(probabilities)}."
+            )
 
-        for i, class_name in enumerate(classes):
-            st.write(f"{class_name.upper()}: {probabilities[i]*100:.2f}%")
+        else:
+            predicted_index = np.argmax(probabilities)
+            predicted_class = classes[predicted_index]
+            confidence = probabilities[predicted_index] * 100
+
+            # -----------------------------
+            # DISPLAY SCORES
+            # -----------------------------
+            st.subheader("📊 Prediction Scores")
+
+            for i, class_name in enumerate(classes):
+                st.write(
+                    f"{class_name.upper()}: {probabilities[i] * 100:.2f}%"
+                )
+
+            # -----------------------------
+            # MAIN RESULT
+            # -----------------------------
+            st.success(
+                f"Predicted Type: {predicted_class.upper()}"
+            )
 
-        st.success(f"Predicted Type: {predicted_class.upper()}")
-        st.info(f"Confidence: {confidence:.2f}%")
+            st.info(
+                f"Confidence: {confidence:.2f}%"
+            )
 
-        # -----------------------------
-        # SUSTAINABILITY TIPS
-        # -----------------------------
-        tips = {
-            'plastic': 'Recycle plastic properly to reduce pollution.',
-            'paper': 'Reuse or recycle paper to save trees.',
-            'metal': 'Metal can be recycled efficiently.',
-            'glass': 'Glass is reusable and recyclable.',
-            'trash': 'Dispose responsibly to reduce environmental damage.',
-            'cardboard': 'Recycle cardboard to reduce waste.'
-        }
-
-        st.subheader("🌱 Sustainability Suggestion")
-        st.write(
-            tips.get(
-                predicted_class,
-                "Dispose responsibly."
+            # -----------------------------
+            # SUSTAINABILITY TIPS
+            # -----------------------------
+            tips = {
+                'plastic': 'Recycle plastic properly to reduce pollution.',
+                'paper': 'Reuse or recycle paper to save trees.',
+                'metal': 'Metal can be recycled efficiently.',
+                'glass': 'Glass is reusable and recyclable.',
+                'trash': 'Dispose responsibly to reduce environmental damage.',
+                'cardboard': 'Recycle cardboard to reduce waste.'
+            }
+
+            st.subheader("🌱 Sustainability Suggestion")
+
+            st.write(
+                tips.get(
+                    predicted_class,
+                    "Dispose responsibly."
+                )
             )
-        )
 
     except UnidentifiedImageError:
         st.error(
@@ -215,10 +241,14 @@ if uploaded_file is not None:
         )
 
     except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
+        st.error(
+            f"Error processing image: {str(e)}"
+        )
 
 # -----------------------------
 # FOOTER
 # -----------------------------
 st.markdown("---")
-st.caption("Built using TensorFlow + Streamlit for Sustainable AI")
+st.caption(
+    "Built using TensorFlow + Streamlit for Sustainable AI"
+)