Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 import tensorflow as tf
 from tensorflow.keras.preprocessing.image import ImageDataGenerator
-from tensorflow.keras.models import Sequential
+from tensorflow.keras.models import Sequential, load_model
 from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
 import numpy as np
 from PIL import Image, UnidentifiedImageError
@@ -12,6 +12,7 @@ import os
 # -----------------------------
 DATASET_DIR = "dataset-resized"
 MODEL_PATH = "waste_classifier.h5"
+CLASS_FILE = "classes.npy"
 IMG_SIZE = (128, 128)
 BATCH_SIZE = 16
 EPOCHS = 5
@@ -25,7 +26,7 @@ st.set_page_config(
 )
 
 # -----------------------------
-# REMOVE CORRUPTED IMAGES
+# CLEAN DATASET
 # -----------------------------
 def clean_dataset(dataset_path):
     valid_extensions = (".jpg", ".jpeg", ".png")
@@ -46,7 +47,7 @@ def clean_dataset(dataset_path):
             try:
                 with Image.open(file_path) as img:
                     img.verify()
-            except (UnidentifiedImageError, OSError):
+            except:
                 try:
                     os.remove(file_path)
                     removed += 1
@@ -105,7 +106,7 @@ def train_model():
         Dense(256, activation='relu'),
         Dropout(0.5),
 
-        Dense(train_data.num_classes, activation='softmax')
+        Dense(len(classes), activation='softmax')
     ])
 
     model.compile(
@@ -114,28 +115,50 @@ def train_model():
         metrics=['accuracy']
     )
 
-    with st.spinner("Training AI model... Please wait."):
+    with st.spinner("Training AI model... Please wait..."):
         model.fit(
             train_data,
             validation_data=val_data,
             epochs=EPOCHS
         )
 
+    # Save model + classes
     model.save(MODEL_PATH)
+    np.save(CLASS_FILE, classes)
 
     return model, classes
 
 # -----------------------------
-# LOAD OR TRAIN MODEL
+# LOAD MODEL
 # -----------------------------
-if not os.path.exists(MODEL_PATH):
-    st.warning("Training model for first-time use. Please wait...")
-    model, classes = train_model()
-else:
-    model = tf.keras.models.load_model(MODEL_PATH)
+def load_or_train_model():
+    if not os.path.exists(MODEL_PATH) or not os.path.exists(CLASS_FILE):
+        st.warning("Training model for first-time use. Please wait...")
+        return train_model()
 
-    # Ensure exact dataset folder order
-    classes = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
+    try:
+        model = load_model(MODEL_PATH)
+        classes = np.load(CLASS_FILE, allow_pickle=True).tolist()
+
+        # Verify output layer
+        output_classes = model.output_shape[-1]
+
+        if output_classes != len(classes):
+            st.warning("Old incorrect model detected. Retraining...")
+            os.remove(MODEL_PATH)
+            os.remove(CLASS_FILE)
+            return train_model()
+
+        return model, classes
+
+    except:
+        st.warning("Model corrupted. Retraining...")
+        return train_model()
+
+# -----------------------------
+# LOAD MODEL
+# -----------------------------
+model, classes = load_or_train_model()
 
 # -----------------------------
 # STREAMLIT UI
@@ -145,15 +168,11 @@ st.write("Upload an image to classify waste and support sustainable recycling.")
 
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
-    type=["jpg", "jpeg", "png"],
-    accept_multiple_files=False
+    type=["jpg", "jpeg", "png"]
 )
 
 if uploaded_file is not None:
     try:
-        # -----------------------------
-        # LOAD IMAGE
-        # -----------------------------
         image = Image.open(uploaded_file).convert("RGB")
 
         st.image(
@@ -163,7 +182,7 @@ if uploaded_file is not None:
         )
 
         # -----------------------------
-        # PREPROCESS IMAGE
+        # PREPROCESS
         # -----------------------------
         img = image.resize(IMG_SIZE)
         img_array = np.array(img) / 255.0
@@ -173,71 +192,55 @@ if uploaded_file is not None:
         # PREDICT
         # -----------------------------
         with st.spinner("Analyzing waste type..."):
-            prediction = model.predict(
-                img_array,
-                verbose=0
-            )
+            prediction = model.predict(img_array, verbose=0)
 
         probabilities = prediction.flatten()
 
+        predicted_index = np.argmax(probabilities)
+        predicted_class = classes[predicted_index]
+        confidence = probabilities[predicted_index] * 100
+
         # -----------------------------
-        # VALIDATE OUTPUT
+        # RESULTS
         # -----------------------------
-        if len(probabilities) != len(classes):
-            st.error(
-                f"Model output mismatch: Expected {len(classes)} classes but got {len(probabilities)}."
-            )
-
-        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.subheader("📊 Prediction Scores")
 
-            st.info(
-                f"Confidence: {confidence:.2f}%"
+        for i, class_name in enumerate(classes):
+            st.write(
+                f"{class_name.upper()}: {probabilities[i]*100:.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.success(
+            f"Predicted Type: {predicted_class.upper()}"
+        )
 
-            st.subheader("🌱 Sustainability Suggestion")
+        st.info(
+            f"Confidence: {confidence:.2f}%"
+        )
 
-            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(
-            "Invalid image file. Please upload a valid JPG, JPEG, or PNG image."
+            "Invalid image file. Please upload a valid image."
         )
 
     except Exception as e: