Update app.py

app.py

@@ -2,7 +2,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.layers import Conv2D, MaxPooling2D, Flatten, Dense
+from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
 import numpy as np
 from PIL import Image, UnidentifiedImageError
 import os
@@ -14,12 +14,15 @@ DATASET_DIR = "dataset-resized"
 MODEL_PATH = "waste_classifier.h5"
 IMG_SIZE = (128, 128)
 BATCH_SIZE = 16
-EPOCHS = 3
+EPOCHS = 5
 
 # -----------------------------
 # PAGE CONFIG
 # -----------------------------
-st.set_page_config(page_title="AI Waste Classifier", layout="centered")
+st.set_page_config(
+    page_title="AI Waste Classifier",
+    layout="centered"
+)
 
 # -----------------------------
 # REMOVE CORRUPTED IMAGES
@@ -61,7 +64,10 @@ def train_model():
 
     datagen = ImageDataGenerator(
         rescale=1./255,
-        validation_split=0.2
+        validation_split=0.2,
+        rotation_range=20,
+        zoom_range=0.2,
+        horizontal_flip=True
     )
 
     train_data = datagen.flow_from_directory(
@@ -69,7 +75,8 @@ def train_model():
         target_size=IMG_SIZE,
         batch_size=BATCH_SIZE,
         class_mode='categorical',
-        subset='training'
+        subset='training',
+        shuffle=True
     )
 
     val_data = datagen.flow_from_directory(
@@ -77,9 +84,12 @@ def train_model():
         target_size=IMG_SIZE,
         batch_size=BATCH_SIZE,
         class_mode='categorical',
-        subset='validation'
+        subset='validation',
+        shuffle=False
     )
 
+    classes = list(train_data.class_indices.keys())
+
     model = Sequential([
         Conv2D(32, (3,3), activation='relu', input_shape=(128,128,3)),
         MaxPooling2D(2,2),
@@ -87,8 +97,13 @@ def train_model():
         Conv2D(64, (3,3), activation='relu'),
         MaxPooling2D(2,2),
 
+        Conv2D(128, (3,3), activation='relu'),
+        MaxPooling2D(2,2),
+
         Flatten(),
-        Dense(128, activation='relu'),
+        Dense(256, activation='relu'),
+        Dropout(0.5),
+
         Dense(train_data.num_classes, activation='softmax')
     ])
 
@@ -107,7 +122,7 @@ def train_model():
 
     model.save(MODEL_PATH)
 
-    return model, list(train_data.class_indices.keys())
+    return model, classes
 
 # -----------------------------
 # LOAD OR TRAIN MODEL
@@ -117,10 +132,13 @@ if not os.path.exists(MODEL_PATH):
     model, classes = train_model()
 else:
     model = tf.keras.models.load_model(MODEL_PATH)
+
+    # IMPORTANT:
+    # Ensure this matches dataset folder order exactly
     classes = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
 # -----------------------------
-# UI
+# STREAMLIT UI
 # -----------------------------
 st.title("♻️ AI Smart Waste Classification")
 st.write("Upload an image to classify waste and support sustainable recycling.")
@@ -141,22 +159,39 @@ if uploaded_file is not None:
             use_container_width=True
         )
 
-        # Preprocess
+        # -----------------------------
+        # PREPROCESS IMAGE
+        # -----------------------------
         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)
+            prediction = model.predict(img_array, verbose=0)
 
-        predicted_class = classes[np.argmax(prediction)]
-        confidence = np.max(prediction) * 100
+        probabilities = prediction[0]
+
+        predicted_index = np.argmax(probabilities)
+        predicted_class = classes[predicted_index]
+        confidence = probabilities[predicted_index] * 100
+
+        # -----------------------------
+        # DISPLAY RESULTS
+        # -----------------------------
+        st.subheader("📊 Prediction Scores")
+
+        for i, class_name in enumerate(classes):
+            st.write(f"{class_name.upper()}: {probabilities[i]*100:.2f}%")
 
-        # Output
         st.success(f"Predicted Type: {predicted_class.upper()}")
         st.info(f"Confidence: {confidence:.2f}%")
 
+        # -----------------------------
+        # SUSTAINABILITY TIPS
+        # -----------------------------
         tips = {
             'plastic': 'Recycle plastic properly to reduce pollution.',
             'paper': 'Reuse or recycle paper to save trees.',
@@ -167,10 +202,17 @@ if uploaded_file is not None:
         }
 
         st.subheader("🌱 Sustainability Suggestion")
-        st.write(tips.get(predicted_class, "Dispose responsibly."))
+        st.write(
+            tips.get(
+                predicted_class,
+                "Dispose responsibly."
+            )
+        )
 
     except UnidentifiedImageError:
-        st.error("Invalid image file. Please upload a proper JPG, JPEG, or PNG 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)}")