Update app.py

app.py

@@ -6,6 +6,7 @@ from tensorflow.keras import layers, models
 import numpy as np
 from PIL import Image, UnidentifiedImageError
 import os
+from sklearn.utils.class_weight import compute_class_weight
 
 # -----------------------------
 # CONFIGURATION
@@ -14,7 +15,7 @@ MODEL_PATH = "waste_classifier.h5"
 DATASET_DIR = "dataset-resized/dataset-resized"
 IMG_SIZE = (128, 128)
 BATCH_SIZE = 32
-EPOCHS = 10
+EPOCHS = 20  # Increased for better accuracy
 
 # Fixed class labels
 CLASSES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
@@ -87,7 +88,10 @@ def train_and_save_model():
 
     datagen = ImageDataGenerator(
         rescale=1./255,
-        validation_split=0.2
+        validation_split=0.2,
+        rotation_range=15,
+        zoom_range=0.1,
+        horizontal_flip=True
     )
 
     train_data = datagen.flow_from_directory(
@@ -110,7 +114,20 @@ def train_and_save_model():
         shuffle=True
     )
 
-    # CNN Model
+    # -----------------------------
+    # CLASS WEIGHTS FOR BALANCED TRAINING
+    # -----------------------------
+    class_weights = compute_class_weight(
+        class_weight='balanced',
+        classes=np.unique(train_data.classes),
+        y=train_data.classes
+    )
+
+    class_weights = dict(enumerate(class_weights))
+
+    # -----------------------------
+    # CNN MODEL
+    # -----------------------------
     model = models.Sequential([
         layers.Input(shape=(128, 128, 3)),
 
@@ -125,9 +142,12 @@ def train_and_save_model():
 
         layers.Flatten(),
 
-        layers.Dense(128, activation='relu'),
+        layers.Dense(256, activation='relu'),
         layers.Dropout(0.5),
 
+        layers.Dense(128, activation='relu'),
+        layers.Dropout(0.3),
+
         layers.Dense(len(CLASSES), activation='softmax')
     ])
 
@@ -144,7 +164,8 @@ def train_and_save_model():
             train_data,
             validation_data=val_data,
             epochs=1,
-            verbose=1
+            verbose=1,
+            class_weight=class_weights
         )
 
         progress_bar.progress((epoch + 1) / EPOCHS)
@@ -208,7 +229,14 @@ def predict_waste(image):
 
     probabilities = prediction[0]
 
-    predicted_index = np.argmax(probabilities)
+    trash_index = CLASSES.index("trash")
+
+    # Trash threshold boost
+    if probabilities[trash_index] > 0.40:
+        predicted_index = trash_index
+    else:
+        predicted_index = np.argmax(probabilities)
+
     predicted_class = CLASSES[predicted_index]
     confidence = probabilities[predicted_index] * 100
 
@@ -285,7 +313,7 @@ if uploaded_file is not None:
         st.subheader("🌱 Sustainability Suggestion")
         st.write(TIPS.get(predicted_class, "Dispose responsibly."))
 
-        # AI Environmental Analysis
+        # AI Analysis
         st.subheader("🤖 AI Environmental Analysis")
         st.success(
             AI_MESSAGES.get(