Update app.py

app.py

@@ -4,7 +4,7 @@ from tensorflow.keras.preprocessing.image import ImageDataGenerator
 from tensorflow.keras.models import Sequential
 from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
 import numpy as np
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 import os
 
 # -----------------------------
@@ -16,10 +16,38 @@ IMG_SIZE = (128, 128)
 BATCH_SIZE = 32
 EPOCHS = 5
 
+# -----------------------------
+# 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
+                continue
+
+            try:
+                with Image.open(file_path) as img:
+                    img.verify()
+            except (UnidentifiedImageError, OSError):
+                os.remove(file_path)
+                removed += 1
+
+    return removed
+
 # -----------------------------
 # TRAIN MODEL FUNCTION
 # -----------------------------
 def train_model():
+    removed_files = clean_dataset(DATASET_DIR)
+    st.info(f"Removed {removed_files} corrupted/invalid files.")
+
     datagen = ImageDataGenerator(
         rescale=1./255,
         validation_split=0.2
@@ -70,7 +98,7 @@ def train_model():
     return model, list(train_data.class_indices.keys())
 
 # -----------------------------
-# LOAD MODEL
+# LOAD OR TRAIN MODEL
 # -----------------------------
 if not os.path.exists(MODEL_PATH):
     st.warning("Training model for first-time use. Please wait...")