Update app.py

app.py

@@ -13,15 +13,19 @@ import os
 DATASET_DIR = "dataset-resized"
 MODEL_PATH = "waste_classifier.h5"
 CLASS_FILE = "classes.npy"
+
 IMG_SIZE = (128, 128)
 BATCH_SIZE = 16
 EPOCHS = 5
 
+# Fixed class labels
+CLASSES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
+
 # -----------------------------
-# PAGE CONFIG
+# PAGE SETTINGS
 # -----------------------------
 st.set_page_config(
-    page_title="AI Waste Classifier",
+    page_title="AI Smart Waste Classification",
     layout="centered"
 )
 
@@ -36,6 +40,7 @@ def clean_dataset(dataset_path):
         for file in files:
             file_path = os.path.join(root, file)
 
+            # Remove invalid extensions
             if not file.lower().endswith(valid_extensions):
                 try:
                     os.remove(file_path)
@@ -44,6 +49,7 @@ def clean_dataset(dataset_path):
                     pass
                 continue
 
+            # Remove corrupted images
             try:
                 with Image.open(file_path) as img:
                     img.verify()
@@ -77,7 +83,8 @@ def train_model():
         batch_size=BATCH_SIZE,
         class_mode='categorical',
         subset='training',
-        shuffle=True
+        shuffle=True,
+        classes=CLASSES
     )
 
     val_data = datagen.flow_from_directory(
@@ -86,11 +93,10 @@ def train_model():
         batch_size=BATCH_SIZE,
         class_mode='categorical',
         subset='validation',
-        shuffle=False
+        shuffle=False,
+        classes=CLASSES
     )
 
-    classes = list(train_data.class_indices.keys())
-
     model = Sequential([
         Conv2D(32, (3,3), activation='relu', input_shape=(128,128,3)),
         MaxPooling2D(2,2),
@@ -106,7 +112,7 @@ def train_model():
         Dense(256, activation='relu'),
         Dropout(0.5),
 
-        Dense(len(classes), activation='softmax')
+        Dense(len(CLASSES), activation='softmax')
     ])
 
     model.compile(
@@ -122,34 +128,31 @@ def train_model():
             epochs=EPOCHS
         )
 
-    # Save model + classes
     model.save(MODEL_PATH)
-    np.save(CLASS_FILE, classes)
+    np.save(CLASS_FILE, CLASSES)
 
-    return model, classes
+    return model
 
 # -----------------------------
-# LOAD MODEL
+# LOAD OR TRAIN
 # -----------------------------
 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...")
+        st.warning("Training model for first-time use...")
         return train_model()
 
     try:
         model = load_model(MODEL_PATH)
-        classes = np.load(CLASS_FILE, allow_pickle=True).tolist()
-
-        # Verify output layer
-        output_classes = model.output_shape[-1]
+        saved_classes = np.load(CLASS_FILE, allow_pickle=True).tolist()
 
-        if output_classes != len(classes):
-            st.warning("Old incorrect model detected. Retraining...")
+        # Force retrain if mismatch
+        if saved_classes != CLASSES or model.output_shape[-1] != len(CLASSES):
+            st.warning("Old model mismatch detected. Retraining...")
             os.remove(MODEL_PATH)
             os.remove(CLASS_FILE)
             return train_model()
 
-        return model, classes
+        return model
 
     except:
         st.warning("Model corrupted. Retraining...")
@@ -158,19 +161,23 @@ def load_or_train_model():
 # -----------------------------
 # LOAD MODEL
 # -----------------------------
-model, classes = load_or_train_model()
+model = load_or_train_model()
 
 # -----------------------------
-# STREAMLIT UI
+# UI
 # -----------------------------
 st.title("♻️ AI Smart Waste Classification")
 st.write("Upload an image to classify waste and support sustainable recycling.")
 
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
-    type=["jpg", "jpeg", "png"]
+    type=["jpg", "jpeg", "png"],
+    accept_multiple_files=False
 )
 
+# -----------------------------
+# PREDICTION
+# -----------------------------
 if uploaded_file is not None:
     try:
         image = Image.open(uploaded_file).convert("RGB")
@@ -181,35 +188,32 @@ if uploaded_file is not None:
             use_container_width=True
         )
 
-        # -----------------------------
-        # PREPROCESS
-        # -----------------------------
+        # Preprocess
         img = image.resize(IMG_SIZE)
         img_array = np.array(img) / 255.0
         img_array = np.expand_dims(img_array, axis=0)
 
-        # -----------------------------
-        # PREDICT
-        # -----------------------------
+        # Predict
         with st.spinner("Analyzing waste type..."):
             prediction = model.predict(img_array, verbose=0)
 
         probabilities = prediction.flatten()
 
         predicted_index = np.argmax(probabilities)
-        predicted_class = classes[predicted_index]
+        predicted_class = CLASSES[predicted_index]
         confidence = probabilities[predicted_index] * 100
 
         # -----------------------------
-        # RESULTS
+        # DISPLAY SCORES
         # -----------------------------
         st.subheader("📊 Prediction Scores")
 
-        for i, class_name in enumerate(classes):
+        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()}"
         )
@@ -218,9 +222,7 @@ if uploaded_file is not None:
             f"Confidence: {confidence:.2f}%"
         )
 
-        # -----------------------------
-        # SUSTAINABILITY TIPS
-        # -----------------------------
+        # Sustainability Tips
         tips = {
             'plastic': 'Recycle plastic properly to reduce pollution.',
             'paper': 'Reuse or recycle paper to save trees.',
@@ -239,19 +241,28 @@ if uploaded_file is not None:
         )
 
     except UnidentifiedImageError:
-        st.error(
-            "Invalid image file. Please upload a valid image."
-        )
+        st.error("Invalid image file. Please upload a valid JPG, JPEG, or PNG image.")
 
     except Exception as e:
-        st.error(
-            f"Error processing image: {str(e)}"
-        )
+        st.error(f"Error processing image: {str(e)}")
+
+# -----------------------------
+# SAMPLE TEST IMAGE IDEAS
+# -----------------------------
+st.markdown("---")
+st.subheader("🖼️ Sample Images to Test")
+st.write("""
+Use images like these:
+- plastic_bottle.jpg
+- newspaper.jpg
+- soda_can.jpg
+- glass_bottle.jpg
+- cardboard_box.jpg
+- trash_bag.jpg
+""")
 
 # -----------------------------
 # FOOTER
 # -----------------------------
 st.markdown("---")
-st.caption(
-    "Built using TensorFlow + Streamlit for Sustainable AI"
-)
+st.caption("Built using TensorFlow + Streamlit for Sustainable AI")