Update app.py (#14)

- Update app.py (1b7c92c619091319bcf24f8d4c9697cd967debb3)

app.py

@@ -18,7 +18,6 @@ IMG_SIZE = (128, 128)
 BATCH_SIZE = 16
 EPOCHS = 5
 
-# Fixed class labels
 CLASSES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
 # -----------------------------
@@ -29,6 +28,32 @@ st.set_page_config(
     layout="centered"
 )
 
+# -----------------------------
+# VALIDATE DATASET
+# -----------------------------
+def validate_dataset():
+    missing_classes = []
+    total_images = 0
+
+    for class_name in CLASSES:
+        class_path = os.path.join(DATASET_DIR, class_name)
+
+        if not os.path.exists(class_path):
+            missing_classes.append(class_name)
+            continue
+
+        files = [
+            f for f in os.listdir(class_path)
+            if f.lower().endswith((".jpg", ".jpeg", ".png"))
+        ]
+
+        total_images += len(files)
+
+        if len(files) == 0:
+            missing_classes.append(class_name)
+
+    return missing_classes, total_images
+
 # -----------------------------
 # CLEAN DATASET
 # -----------------------------
@@ -40,7 +65,6 @@ def clean_dataset(dataset_path):
         for file in files:
             file_path = os.path.join(root, file)
 
-            # Remove invalid extensions
             if not file.lower().endswith(valid_extensions):
                 try:
                     os.remove(file_path)
@@ -49,7 +73,6 @@ def clean_dataset(dataset_path):
                     pass
                 continue
 
-            # Remove corrupted images
             try:
                 with Image.open(file_path) as img:
                     img.verify()
@@ -66,6 +89,18 @@ def clean_dataset(dataset_path):
 # TRAIN MODEL
 # -----------------------------
 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.")
 
@@ -97,6 +132,13 @@ def train_model():
         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),
@@ -134,7 +176,7 @@ def train_model():
     return model
 
 # -----------------------------
-# LOAD OR TRAIN
+# LOAD OR TRAIN MODEL
 # -----------------------------
 def load_or_train_model():
     if not os.path.exists(MODEL_PATH) or not os.path.exists(CLASS_FILE):
@@ -145,7 +187,6 @@ def load_or_train_model():
         model = load_model(MODEL_PATH)
         saved_classes = np.load(CLASS_FILE, allow_pickle=True).tolist()
 
-        # Force retrain if mismatch
         if saved_classes != CLASSES or model.output_shape[-1] != len(CLASSES):
             st.warning("Old model mismatch detected. Retraining...")
             os.remove(MODEL_PATH)
@@ -171,13 +212,9 @@ st.write("Upload an image to classify waste and support sustainable recycling.")
 
 uploaded_file = st.file_uploader(
     "Upload Waste Image",
-    type=["jpg", "jpeg", "png"],
-    accept_multiple_files=False
+    type=["jpg", "jpeg", "png"]
 )
 
-# -----------------------------
-# PREDICTION
-# -----------------------------
 if uploaded_file is not None:
     try:
         image = Image.open(uploaded_file).convert("RGB")
@@ -188,12 +225,10 @@ if uploaded_file is not None:
             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, verbose=0)
 
@@ -203,9 +238,6 @@ if uploaded_file is not None:
         predicted_class = CLASSES[predicted_index]
         confidence = probabilities[predicted_index] * 100
 
-        # -----------------------------
-        # DISPLAY SCORES
-        # -----------------------------
         st.subheader("📊 Prediction Scores")
 
         for i, class_name in enumerate(CLASSES):
@@ -213,7 +245,6 @@ if uploaded_file is not None:
                 f"{class_name.upper()}: {probabilities[i]*100:.2f}%"
             )
 
-        # Main result
         st.success(
             f"Predicted Type: {predicted_class.upper()}"
         )
@@ -222,7 +253,6 @@ if uploaded_file is not None:
             f"Confidence: {confidence:.2f}%"
         )
 
-        # Sustainability Tips
         tips = {
             'plastic': 'Recycle plastic properly to reduce pollution.',
             'paper': 'Reuse or recycle paper to save trees.',
@@ -241,26 +271,11 @@ if uploaded_file is not None:
         )
 
     except UnidentifiedImageError:
-        st.error("Invalid image file. Please upload a valid JPG, JPEG, or PNG image.")
+        st.error("Invalid image file uploaded.")
 
     except Exception as e:
         st.error(f"Error processing image: {str(e)}")
 
-# -----------------------------
-# SAMPLE TEST IMAGE IDEAS
-# -----------------------------
-st.markdown("---")
-st.subheader("🖼️ Sample Images to Test")
-st.write("""
-Use images like these:
-- plastic_bottle.jpg
-- newspaper.jpg
-- soda_can.jpg
-- glass_bottle.jpg
-- cardboard_box.jpg
-- trash_bag.jpg
-""")
-
 # -----------------------------
 # FOOTER
 # -----------------------------