Update app.py

app.py

@@ -38,13 +38,44 @@ st.set_page_config(
 )
 
 # -----------------------------
-# TRAIN MODEL
+# VALIDATE DATASET STRUCTURE
 # -----------------------------
-def train_and_save_model():
+def validate_dataset():
+    """
+    Ensure dataset folder contains 6 required class subfolders
+    """
     if not os.path.exists(DATASET_DIR):
         st.error(f"❌ Dataset folder '{DATASET_DIR}' not found.")
         st.stop()
 
+    found_folders = sorted([
+        folder for folder in os.listdir(DATASET_DIR)
+        if os.path.isdir(os.path.join(DATASET_DIR, folder))
+    ])
+
+    missing = [cls for cls in CLASSES if cls not in found_folders]
+
+    if missing:
+        st.error("❌ Dataset structure is incorrect.")
+        st.write("Expected folders:")
+        for cls in CLASSES:
+            st.write(f"- {cls}")
+
+        st.write("Missing folders:")
+        for m in missing:
+            st.write(f"- {m}")
+
+        st.stop()
+
+    return found_folders
+
+
+# -----------------------------
+# TRAIN MODEL
+# -----------------------------
+def train_and_save_model():
+    validate_dataset()
+
     st.info("⚙️ Model not found. Training a new model... This may take several minutes.")
 
     datagen = ImageDataGenerator(
@@ -52,12 +83,11 @@ 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,
+        classes=CLASSES,  # FORCE CORRECT CLASSES
         class_mode='categorical',
         subset='training',
         shuffle=True
@@ -67,18 +97,12 @@ def train_and_save_model():
         DATASET_DIR,
         target_size=IMG_SIZE,
         batch_size=BATCH_SIZE,
+        classes=CLASSES,
         class_mode='categorical',
         subset='validation',
         shuffle=True
     )
 
-    # 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)),
@@ -100,7 +124,6 @@ def train_and_save_model():
         layers.Dense(len(CLASSES), activation='softmax')
     ])
 
-    # COMPILE FIX:
     model.compile(
         optimizer='adam',
         loss='categorical_crossentropy',
@@ -178,9 +201,6 @@ 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
@@ -201,18 +221,25 @@ with st.sidebar:
     st.header("📂 Dataset Status")
 
     if os.path.exists(DATASET_DIR):
-        st.success("Dataset Found")
+        folders = sorted([
+            folder for folder in os.listdir(DATASET_DIR)
+            if os.path.isdir(os.path.join(DATASET_DIR, folder))
+        ])
+
+        if folders:
+            st.success("Dataset Found")
 
-        folders = sorted(os.listdir(DATASET_DIR))
+            for folder in folders:
+                st.write(f"✔️ {folder}")
 
-        for folder in folders:
-            st.write(f"✔️ {folder}")
+        else:
+            st.error("No class folders found")
 
     else:
         st.error("Dataset Missing")
 
 # -----------------------------
-# FILE UPLOAD
+# FILE UPLOADER
 # -----------------------------
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
@@ -235,7 +262,6 @@ if uploaded_file is not None:
         with st.spinner("🔍 Analyzing waste type..."):
             predicted_class, confidence, probabilities = predict_waste(image)
 
-        # Results
         st.subheader("📊 Prediction Scores")
 
         for i, class_name in enumerate(CLASSES):
@@ -246,7 +272,6 @@ if uploaded_file is not None:
         st.info(f"🎯 Confidence: {confidence:.2f}%")
         st.write(f"📁 Uploaded File: {uploaded_file.name}")
 
-        # Sustainability tip
         st.subheader("🌱 Sustainability Suggestion")
         st.write(TIPS.get(predicted_class, "Dispose responsibly."))
 
@@ -262,13 +287,15 @@ if uploaded_file is not None:
 st.markdown("---")
 st.subheader("🖼️ Recommended Test Images")
 st.write("""
-Try uploading:
-- plastic_bottle.jpg
-- glass_bottle.jpg
-- cardboard_box.jpg
-- soda_can.jpg
-- newspaper.jpg
-- trash_bag.jpg
+Your dataset folder should look like:
+
+dataset-resized/
+├── cardboard/
+├── glass/
+├── metal/
+├── paper/
+├── plastic/
+└── trash/
 """)
 
 # -----------------------------