Update app.py

app.py

@@ -13,36 +13,47 @@ import os
 DATASET_DIR = "dataset-resized"
 MODEL_PATH = "waste_classifier.h5"
 IMG_SIZE = (128, 128)
-BATCH_SIZE = 32
-EPOCHS = 5
+BATCH_SIZE = 16
+EPOCHS = 3
+
+# -----------------------------
+# PAGE CONFIG
+# -----------------------------
+st.set_page_config(page_title="AI Waste Classifier", layout="centered")
 
 # -----------------------------
 # REMOVE CORRUPTED IMAGES
 # -----------------------------
 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):
-                os.remove(file_path)
-                removed += 1
+                try:
+                    os.remove(file_path)
+                    removed += 1
+                except:
+                    pass
                 continue
 
             try:
                 with Image.open(file_path) as img:
                     img.verify()
             except (UnidentifiedImageError, OSError):
-                os.remove(file_path)
-                removed += 1
+                try:
+                    os.remove(file_path)
+                    removed += 1
+                except:
+                    pass
 
     return removed
 
 # -----------------------------
-# TRAIN MODEL FUNCTION
+# TRAIN MODEL
 # -----------------------------
 def train_model():
     removed_files = clean_dataset(DATASET_DIR)
@@ -87,11 +98,12 @@ def train_model():
         metrics=['accuracy']
     )
 
-    model.fit(
-        train_data,
-        validation_data=val_data,
-        epochs=EPOCHS
-    )
+    with st.spinner("Training AI model... Please wait."):
+        model.fit(
+            train_data,
+            validation_data=val_data,
+            epochs=EPOCHS
+        )
 
     model.save(MODEL_PATH)
 
@@ -108,47 +120,60 @@ else:
     classes = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
 # -----------------------------
-# STREAMLIT UI
+# UI
 # -----------------------------
-st.set_page_config(page_title="AI Waste Classifier")
 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
 )
 
 if uploaded_file is not None:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    # Preprocess image
-    img = image.resize(IMG_SIZE)
-    img_array = np.array(img) / 255.0
-    img_array = np.expand_dims(img_array, axis=0)
-
-    # Predict
-    prediction = model.predict(img_array)
-    predicted_class = classes[np.argmax(prediction)]
-    confidence = np.max(prediction) * 100
-
-    # Display Results
-    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.',
-        '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.subheader("🌱 Sustainability Suggestion")
-    st.write(tips.get(predicted_class, "Dispose responsibly."))
+    try:
+        image = Image.open(uploaded_file).convert("RGB")
+
+        st.image(
+            image,
+            caption="Uploaded Image",
+            use_container_width=True
+        )
+
+        # Preprocess
+        img = image.resize(IMG_SIZE)
+        img_array = np.array(img) / 255.0
+        img_array = np.expand_dims(img_array, axis=0)
+
+        # Predict
+        with st.spinner("Analyzing waste type..."):
+            prediction = model.predict(img_array)
+
+        predicted_class = classes[np.argmax(prediction)]
+        confidence = np.max(prediction) * 100
+
+        # Output
+        st.success(f"Predicted Type: {predicted_class.upper()}")
+        st.info(f"Confidence: {confidence:.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.subheader("🌱 Sustainability Suggestion")
+        st.write(tips.get(predicted_class, "Dispose responsibly."))
+
+    except UnidentifiedImageError:
+        st.error("Invalid image file. Please upload a proper JPG, JPEG, or PNG image.")
+
+    except Exception as e:
+        st.error(f"Error processing image: {str(e)}")
 
 # -----------------------------
 # FOOTER