Update app.py

app.py

@@ -38,12 +38,9 @@ st.set_page_config(
 )
 
 # -----------------------------
-# MODEL TRAINING FUNCTION
+# TRAIN MODEL
 # -----------------------------
 def train_and_save_model():
-    """
-    Train CNN model if model file doesn't exist
-    """
     if not os.path.exists(DATASET_DIR):
         st.error(f"❌ Dataset folder '{DATASET_DIR}' not found.")
         st.stop()
@@ -55,12 +52,15 @@ def train_and_save_model():
         validation_split=0.2
     )
 
+    # IMPORTANT FIX:
+    # Use categorical labels instead of binary
     train_data = datagen.flow_from_directory(
         DATASET_DIR,
         target_size=IMG_SIZE,
         batch_size=BATCH_SIZE,
         class_mode='categorical',
-        subset='training'
+        subset='training',
+        shuffle=True
     )
 
     val_data = datagen.flow_from_directory(
@@ -68,21 +68,29 @@ def train_and_save_model():
         target_size=IMG_SIZE,
         batch_size=BATCH_SIZE,
         class_mode='categorical',
-        subset='validation'
+        subset='validation',
+        shuffle=True
     )
 
-    # CNN Architecture
+    # Verify class count
+    if train_data.num_classes != len(CLASSES):
+        st.error(
+            f"❌ Dataset class mismatch. Expected {len(CLASSES)} classes but found {train_data.num_classes}."
+        )
+        st.stop()
+
+    # CNN Model
     model = models.Sequential([
-        layers.Input(shape=(128,128,3)),
+        layers.Input(shape=(128, 128, 3)),
 
-        layers.Conv2D(32, (3,3), activation='relu'),
-        layers.MaxPooling2D(2,2),
+        layers.Conv2D(32, (3, 3), activation='relu'),
+        layers.MaxPooling2D(2, 2),
 
-        layers.Conv2D(64, (3,3), activation='relu'),
-        layers.MaxPooling2D(2,2),
+        layers.Conv2D(64, (3, 3), activation='relu'),
+        layers.MaxPooling2D(2, 2),
 
-        layers.Conv2D(128, (3,3), activation='relu'),
-        layers.MaxPooling2D(2,2),
+        layers.Conv2D(128, (3, 3), activation='relu'),
+        layers.MaxPooling2D(2, 2),
 
         layers.Flatten(),
 
@@ -92,13 +100,13 @@ def train_and_save_model():
         layers.Dense(len(CLASSES), activation='softmax')
     ])
 
+    # COMPILE FIX:
     model.compile(
         optimizer='adam',
         loss='categorical_crossentropy',
         metrics=['accuracy']
     )
 
-    # Progress bar
     progress_bar = st.progress(0)
 
     for epoch in range(EPOCHS):
@@ -106,8 +114,9 @@ def train_and_save_model():
             train_data,
             validation_data=val_data,
             epochs=1,
-            verbose=0
+            verbose=1
         )
+
         progress_bar.progress((epoch + 1) / EPOCHS)
 
     model.save(MODEL_PATH)
@@ -118,7 +127,7 @@ def train_and_save_model():
 
 
 # -----------------------------
-# LOAD OR TRAIN MODEL
+# LOAD MODEL
 # -----------------------------
 @st.cache_resource
 def load_ai_model():
@@ -127,13 +136,13 @@ def load_ai_model():
             model = load_model(MODEL_PATH)
 
             if model.output_shape[-1] != len(CLASSES):
-                st.warning("⚠️ Model output mismatch. Retraining model...")
+                st.warning("⚠️ Model mismatch. Retraining...")
                 return train_and_save_model()
 
             return model
 
         except Exception:
-            st.warning("⚠️ Corrupted model file. Retraining...")
+            st.warning("⚠️ Corrupted model. Retraining...")
             return train_and_save_model()
 
     else:
@@ -143,7 +152,7 @@ def load_ai_model():
 model = load_ai_model()
 
 # -----------------------------
-# IMAGE PREPROCESSING
+# PREPROCESS IMAGE
 # -----------------------------
 def preprocess_image(image):
     image = image.convert("RGB")
@@ -160,7 +169,7 @@ def preprocess_image(image):
 
 
 # -----------------------------
-# PREDICTION FUNCTION
+# PREDICT
 # -----------------------------
 def predict_waste(image):
     processed_img = preprocess_image(image)
@@ -169,6 +178,9 @@ def predict_waste(image):
 
     probabilities = prediction[0]
 
+    if len(probabilities) != len(CLASSES):
+        raise ValueError("Prediction output mismatch.")
+
     predicted_index = np.argmax(probabilities)
     predicted_class = CLASSES[predicted_index]
     confidence = probabilities[predicted_index] * 100
@@ -180,10 +192,10 @@ def predict_waste(image):
 # UI HEADER
 # -----------------------------
 st.title("♻️ AI Smart Waste Classification")
-st.write("Upload an image to classify waste and encourage sustainable recycling.")
+st.write("Upload an image to classify waste and support sustainable recycling.")
 
 # -----------------------------
-# DATASET CHECK
+# SIDEBAR
 # -----------------------------
 with st.sidebar:
     st.header("📂 Dataset Status")
@@ -191,14 +203,16 @@ with st.sidebar:
     if os.path.exists(DATASET_DIR):
         st.success("Dataset Found")
 
-        for folder in os.listdir(DATASET_DIR):
+        folders = sorted(os.listdir(DATASET_DIR))
+
+        for folder in folders:
             st.write(f"✔️ {folder}")
 
     else:
         st.error("Dataset Missing")
 
 # -----------------------------
-# FILE UPLOADER
+# FILE UPLOAD
 # -----------------------------
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
@@ -206,7 +220,7 @@ uploaded_file = st.file_uploader(
 )
 
 # -----------------------------
-# IMAGE ANALYSIS
+# ANALYSIS
 # -----------------------------
 if uploaded_file is not None:
     try:
@@ -226,18 +240,18 @@ if uploaded_file is not None:
 
         for i, class_name in enumerate(CLASSES):
             st.progress(float(probabilities[i]))
-            st.write(f"{class_name.upper()}: {probabilities[i]*100:.2f}%")
+            st.write(f"{class_name.upper()}: {probabilities[i] * 100:.2f}%")
 
         st.success(f"✅ Predicted Type: {predicted_class.upper()}")
         st.info(f"🎯 Confidence: {confidence:.2f}%")
         st.write(f"📁 Uploaded File: {uploaded_file.name}")
 
-        # Sustainability Tip
+        # Sustainability tip
         st.subheader("🌱 Sustainability Suggestion")
         st.write(TIPS.get(predicted_class, "Dispose responsibly."))
 
     except UnidentifiedImageError:
-        st.error("❌ Invalid image format. Upload JPG, JPEG, or PNG.")
+        st.error("❌ Invalid image file. Upload JPG, JPEG, or PNG.")
 
     except Exception as e:
         st.error(f"❌ Error processing image: {str(e)}")