Update app.py

app.py

@@ -1,8 +1,6 @@
 import streamlit as st
 import tensorflow as tf
-from tensorflow.keras.preprocessing.image import ImageDataGenerator
-from tensorflow.keras.models import Sequential, load_model
-from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
+from tensorflow.keras.models import load_model
 import numpy as np
 from PIL import Image, UnidentifiedImageError
 import os
@@ -10,16 +8,22 @@ import os
 # -----------------------------
 # CONFIGURATION
 # -----------------------------
-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']
 
+# 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.'
+}
+
 # -----------------------------
 # PAGE SETTINGS
 # -----------------------------
@@ -29,252 +33,150 @@ st.set_page_config(
 )
 
 # -----------------------------
-# VALIDATE DATASET
+# LOAD MODEL
 # -----------------------------
-def validate_dataset():
-    missing_classes = []
-    total_images = 0
-
-    for class_name in CLASSES:
-        class_path = os.path.join(DATASET_DIR, class_name)
+@st.cache_resource
+def load_ai_model():
+    """
+    Load trained TensorFlow model safely
+    """
+    if not os.path.exists(MODEL_PATH):
+        st.error("❌ Model file 'waste_classifier.h5' not found.")
+        st.stop()
 
-        if not os.path.exists(class_path):
-            missing_classes.append(class_name)
-            continue
+    try:
+        model = load_model(MODEL_PATH)
 
-        files = [
-            f for f in os.listdir(class_path)
-            if f.lower().endswith((".jpg", ".jpeg", ".png"))
-        ]
+        # 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()
 
-        total_images += len(files)
+        return model
 
-        if len(files) == 0:
-            missing_classes.append(class_name)
+    except Exception as e:
+        st.error(f"❌ Error loading model: {str(e)}")
+        st.stop()
 
-    return missing_classes, total_images
 
-# -----------------------------
-# CLEAN DATASET
-# -----------------------------
-def clean_dataset(dataset_path):
-    valid_extensions = (".jpg", ".jpeg", ".png")
-    removed = 0
-
-    for root, dirs, files in os.walk(dataset_path):
-        for file in files:
-            file_path = os.path.join(root, file)
-
-            if not file.lower().endswith(valid_extensions):
-                try:
-                    os.remove(file_path)
-                    removed += 1
-                except:
-                    pass
-                continue
-
-            try:
-                with Image.open(file_path) as img:
-                    img.verify()
-            except:
-                try:
-                    os.remove(file_path)
-                    removed += 1
-                except:
-                    pass
-
-    return removed
+model = load_ai_model()
 
 # -----------------------------
-# TRAIN MODEL
+# IMAGE PREPROCESSING FUNCTION
 # -----------------------------
-def train_model():
-    missing_classes, total_images = validate_dataset()
-
-    if total_images == 0:
-        st.error("Dataset is empty. Please upload proper waste images.")
-        st.stop()
-
-    if missing_classes:
-        st.error(
-            f"Missing or empty class folders: {', '.join(missing_classes)}"
-        )
-        st.stop()
-
-    removed_files = clean_dataset(DATASET_DIR)
-    st.info(f"Removed {removed_files} corrupted/invalid files.")
-
-    datagen = ImageDataGenerator(
-        rescale=1./255,
-        validation_split=0.2,
-        rotation_range=20,
-        zoom_range=0.2,
-        horizontal_flip=True
-    )
-
-    train_data = datagen.flow_from_directory(
-        DATASET_DIR,
-        target_size=IMG_SIZE,
-        batch_size=BATCH_SIZE,
-        class_mode='categorical',
-        subset='training',
-        shuffle=True,
-        classes=CLASSES
-    )
-
-    val_data = datagen.flow_from_directory(
-        DATASET_DIR,
-        target_size=IMG_SIZE,
-        batch_size=BATCH_SIZE,
-        class_mode='categorical',
-        subset='validation',
-        shuffle=False,
-        classes=CLASSES
-    )
-
-    # Safety check
-    if train_data.samples == 0 or val_data.samples == 0:
-        st.error(
-            "Dataset loading failed. Ensure each folder contains enough valid images."
-        )
-        st.stop()
-
-    model = Sequential([
-        Conv2D(32, (3,3), activation='relu', input_shape=(128,128,3)),
-        MaxPooling2D(2,2),
-
-        Conv2D(64, (3,3), activation='relu'),
-        MaxPooling2D(2,2),
-
-        Conv2D(128, (3,3), activation='relu'),
-        MaxPooling2D(2,2),
-
-        Flatten(),
-
-        Dense(256, activation='relu'),
-        Dropout(0.5),
+def preprocess_image(image):
+    """
+    Resize and normalize uploaded image
+    """
+    image = image.convert("RGB")
+    image = image.resize(IMG_SIZE)
 
-        Dense(len(CLASSES), activation='softmax')
-    ])
+    img_array = np.array(image, dtype=np.float32) / 255.0
 
-    model.compile(
-        optimizer='adam',
-        loss='categorical_crossentropy',
-        metrics=['accuracy']
-    )
+    # Ensure proper shape
+    if img_array.shape != (128, 128, 3):
+        raise ValueError("Image shape mismatch after preprocessing.")
 
-    with st.spinner("Training AI model... Please wait..."):
-        model.fit(
-            train_data,
-            validation_data=val_data,
-            epochs=EPOCHS
-        )
+    img_array = np.expand_dims(img_array, axis=0)
 
-    model.save(MODEL_PATH)
-    np.save(CLASS_FILE, CLASSES)
+    return img_array
 
-    return model
 
 # -----------------------------
-# LOAD OR TRAIN MODEL
+# PREDICTION FUNCTION
 # -----------------------------
-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...")
-        return train_model()
+def predict_waste(image):
+    """
+    Predict waste category
+    """
+    processed_img = preprocess_image(image)
 
-    try:
-        model = load_model(MODEL_PATH)
-        saved_classes = np.load(CLASS_FILE, allow_pickle=True).tolist()
+    prediction = model.predict(processed_img, verbose=0)
 
-        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()
+    probabilities = prediction[0]
 
-        return model
+    if len(probabilities) != len(CLASSES):
+        raise ValueError("Prediction output size mismatch.")
 
-    except:
-        st.warning("Model corrupted. Retraining...")
-        return train_model()
+    predicted_index = np.argmax(probabilities)
+    predicted_class = CLASSES[predicted_index]
+    confidence = probabilities[predicted_index] * 100
+
+    return predicted_class, confidence, probabilities
 
-# -----------------------------
-# LOAD MODEL
-# -----------------------------
-model = load_or_train_model()
 
 # -----------------------------
-# UI
+# UI HEADER
 # -----------------------------
 st.title("♻️ AI Smart Waste Classification")
-st.write("Upload an image to classify waste and support sustainable recycling.")
+st.write("Upload an image to classify waste for smart recycling.")
 
+# -----------------------------
+# FILE UPLOAD
+# -----------------------------
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
     type=["jpg", "jpeg", "png"]
 )
 
+# -----------------------------
+# IMAGE PREDICTION
+# -----------------------------
 if uploaded_file is not None:
     try:
-        image = Image.open(uploaded_file).convert("RGB")
+        # Load image
+        image = Image.open(uploaded_file)
 
+        # Display image
         st.image(
             image,
-            caption="Uploaded Image",
+            caption=f"Uploaded Image: {uploaded_file.name}",
             use_container_width=True
         )
 
-        img = image.resize(IMG_SIZE)
-        img_array = np.array(img) / 255.0
-        img_array = np.expand_dims(img_array, axis=0)
-
-        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]
-        confidence = probabilities[predicted_index] * 100
+        # Predict
+        with st.spinner("🔍 Analyzing waste type..."):
+            predicted_class, confidence, probabilities = predict_waste(image)
 
+        # -----------------------------
+        # DISPLAY RESULTS
+        # -----------------------------
         st.subheader("📊 Prediction Scores")
 
         for i, class_name in enumerate(CLASSES):
-            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.progress(float(probabilities[i]))
+            st.write(f"{class_name.upper()}: {probabilities[i] * 100:.2f}%")
 
-        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.info(f"🎯 Confidence: {confidence:.2f}%")
+        st.write(f"📁 Uploaded File: {uploaded_file.name}")
 
+        # Sustainability tip
         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 uploaded.")
+        st.error("❌ Invalid image file. Please upload JPG, JPEG, or PNG.")
 
     except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
+        st.error(f"❌ Error processing image: {str(e)}")
+
+# -----------------------------
+# SAMPLE GUIDE
+# -----------------------------
+st.markdown("---")
+st.subheader("🖼️ Recommended Test Images")
+st.write("""
+Try uploading:
+- plastic_bottle.jpg
+- glass_bottle.jpg
+- cardboard_box.jpg
+- soda_can.jpg
+- newspaper.jpg
+- trash_bag.jpg
+""")
 
 # -----------------------------
 # FOOTER