Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 import tensorflow as tf
 from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras import layers, models
 import numpy as np
 from PIL import Image, UnidentifiedImageError
 import os
@@ -9,7 +11,10 @@ import os
 # CONFIGURATION
 # -----------------------------
 MODEL_PATH = "waste_classifier.h5"
+DATASET_DIR = "dataset-resized"
 IMG_SIZE = (128, 128)
+BATCH_SIZE = 32
+EPOCHS = 10
 
 # Fixed class labels
 CLASSES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
@@ -33,51 +38,121 @@ st.set_page_config(
 )
 
 # -----------------------------
-# LOAD MODEL
+# MODEL TRAINING FUNCTION
 # -----------------------------
-@st.cache_resource
-def load_ai_model():
+def train_and_save_model():
     """
-    Load trained TensorFlow model safely
+    Train CNN model if model file doesn't exist
     """
-    if not os.path.exists(MODEL_PATH):
-        st.error("❌ Model file 'waste_classifier.h5' not found.")
+    if not os.path.exists(DATASET_DIR):
+        st.error(f"❌ Dataset folder '{DATASET_DIR}' not found.")
         st.stop()
 
-    try:
-        model = load_model(MODEL_PATH)
+    st.info("⚙️ Model not found. Training a new model... This may take several minutes.")
+
+    datagen = ImageDataGenerator(
+        rescale=1./255,
+        validation_split=0.2
+    )
+
+    train_data = datagen.flow_from_directory(
+        DATASET_DIR,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode='categorical',
+        subset='training'
+    )
+
+    val_data = datagen.flow_from_directory(
+        DATASET_DIR,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode='categorical',
+        subset='validation'
+    )
+
+    # CNN Architecture
+    model = models.Sequential([
+        layers.Input(shape=(128,128,3)),
+
+        layers.Conv2D(32, (3,3), activation='relu'),
+        layers.MaxPooling2D(2,2),
+
+        layers.Conv2D(64, (3,3), activation='relu'),
+        layers.MaxPooling2D(2,2),
+
+        layers.Conv2D(128, (3,3), activation='relu'),
+        layers.MaxPooling2D(2,2),
+
+        layers.Flatten(),
+
+        layers.Dense(128, activation='relu'),
+        layers.Dropout(0.5),
+
+        layers.Dense(len(CLASSES), activation='softmax')
+    ])
+
+    model.compile(
+        optimizer='adam',
+        loss='categorical_crossentropy',
+        metrics=['accuracy']
+    )
+
+    # Progress bar
+    progress_bar = st.progress(0)
+
+    for epoch in range(EPOCHS):
+        model.fit(
+            train_data,
+            validation_data=val_data,
+            epochs=1,
+            verbose=0
+        )
+        progress_bar.progress((epoch + 1) / EPOCHS)
 
-        # Validate output classes
-        if model.output_shape[-1] != len(CLASSES):
-            st.error(
-                f"❌ Model output mismatch. Expected {len(CLASSES)} classes, got {model.output_shape[-1]}."
-            )
-            st.stop()
+    model.save(MODEL_PATH)
 
-        return model
+    st.success("✅ Model trained and saved successfully!")
 
-    except Exception as e:
-        st.error(f"❌ Error loading model: {str(e)}")
-        st.stop()
+    return model
+
+
+# -----------------------------
+# LOAD OR TRAIN MODEL
+# -----------------------------
+@st.cache_resource
+def load_ai_model():
+    if os.path.exists(MODEL_PATH):
+        try:
+            model = load_model(MODEL_PATH)
+
+            if model.output_shape[-1] != len(CLASSES):
+                st.warning("⚠️ Model output mismatch. Retraining model...")
+                return train_and_save_model()
+
+            return model
+
+        except Exception:
+            st.warning("⚠️ Corrupted model file. Retraining...")
+            return train_and_save_model()
+
+    else:
+        return train_and_save_model()
 
 
 model = load_ai_model()
 
 # -----------------------------
-# IMAGE PREPROCESSING FUNCTION
+# IMAGE PREPROCESSING
 # -----------------------------
 def preprocess_image(image):
-    """
-    Resize and normalize uploaded image
-    """
     image = image.convert("RGB")
     image = image.resize(IMG_SIZE)
 
     img_array = np.array(image, dtype=np.float32) / 255.0
 
-    # Ensure proper shape
     if img_array.shape != (128, 128, 3):
-        raise ValueError("Image shape mismatch after preprocessing.")
+        raise ValueError("Image preprocessing failed.")
 
     img_array = np.expand_dims(img_array, axis=0)
 
@@ -88,18 +163,12 @@ def preprocess_image(image):
 # PREDICTION FUNCTION
 # -----------------------------
 def predict_waste(image):
-    """
-    Predict waste category
-    """
     processed_img = preprocess_image(image)
 
     prediction = model.predict(processed_img, verbose=0)
 
     probabilities = prediction[0]
 
-    if len(probabilities) != len(CLASSES):
-        raise ValueError("Prediction output size mismatch.")
-
     predicted_index = np.argmax(probabilities)
     predicted_class = CLASSES[predicted_index]
     confidence = probabilities[predicted_index] * 100
@@ -111,10 +180,25 @@ def predict_waste(image):
 # UI HEADER
 # -----------------------------
 st.title("♻️ AI Smart Waste Classification")
-st.write("Upload an image to classify waste for smart recycling.")
+st.write("Upload an image to classify waste and encourage sustainable recycling.")
+
+# -----------------------------
+# DATASET CHECK
+# -----------------------------
+with st.sidebar:
+    st.header("📂 Dataset Status")
+
+    if os.path.exists(DATASET_DIR):
+        st.success("Dataset Found")
+
+        for folder in os.listdir(DATASET_DIR):
+            st.write(f"✔️ {folder}")
+
+    else:
+        st.error("Dataset Missing")
 
 # -----------------------------
-# FILE UPLOAD
+# FILE UPLOADER
 # -----------------------------
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
@@ -122,43 +206,38 @@ uploaded_file = st.file_uploader(
 )
 
 # -----------------------------
-# IMAGE PREDICTION
+# IMAGE ANALYSIS
 # -----------------------------
 if uploaded_file is not None:
     try:
-        # Load image
         image = Image.open(uploaded_file)
 
-        # Display image
         st.image(
             image,
             caption=f"Uploaded Image: {uploaded_file.name}",
             use_container_width=True
         )
 
-        # Predict
         with st.spinner("🔍 Analyzing waste type..."):
             predicted_class, confidence, probabilities = predict_waste(image)
 
-        # -----------------------------
-        # DISPLAY RESULTS
-        # -----------------------------
+        # Results
         st.subheader("📊 Prediction Scores")
 
         for i, class_name in enumerate(CLASSES):
             st.progress(float(probabilities[i]))
-            st.write(f"{class_name.upper()}: {probabilities[i] * 100:.2f}%")
+            st.write(f"{class_name.upper()}: {probabilities[i]*100:.2f}%")
 
         st.success(f"✅ Predicted Type: {predicted_class.upper()}")
         st.info(f"🎯 Confidence: {confidence:.2f}%")
         st.write(f"📁 Uploaded File: {uploaded_file.name}")
 
-        # Sustainability tip
+        # Sustainability Tip
         st.subheader("🌱 Sustainability Suggestion")
         st.write(TIPS.get(predicted_class, "Dispose responsibly."))
 
     except UnidentifiedImageError:
-        st.error("❌ Invalid image file. Please upload JPG, JPEG, or PNG.")
+        st.error("❌ Invalid image format. Upload JPG, JPEG, or PNG.")
 
     except Exception as e:
         st.error(f"❌ Error processing image: {str(e)}")