Update app.py

app.py

@@ -1,24 +1,50 @@
-
 import tensorflow as tf
 import os
-# Avoid OOM errors by setting GPU Memory Consumption Growth
-gpus = tf.config.experimental.list_physical_devices('GPU')
-for gpu in gpus: 
-    tf.config.experimental.set_memory_growth(gpu, True)
-tf.config.list_physical_devices('GPU')
 import cv2
 import imghdr
-data_dir = 'data' 
-image_exts = ['jpeg','jpg', 'bmp', 'png']
+from PIL import Image
+
+# Configure GPU (if available)
+gpus = tf.config.list_physical_devices('GPU')
+if gpus:
+    for gpu in gpus:
+        tf.config.experimental.set_memory_growth(gpu, True)
+
+print("Available GPUs:", tf.config.list_physical_devices('GPU'))
+
+# Allowed image extensions
+image_exts = ['jpeg', 'jpg', 'bmp', 'png']
+
+# Root data directory
+data_dir = 'data'
+
+# Walk through directory and clean up images
 for image_class in os.listdir(data_dir): 
-    for image in os.listdir(os.path.join(data_dir, image_class)):
-        image_path = os.path.join(data_dir, image_class, image)
+    class_path = os.path.join(data_dir, image_class)
+    for image in os.listdir(class_path):
+        image_path = os.path.join(class_path, image)
         try: 
             img = cv2.imread(image_path)
+            if img is None:
+                print('Unreadable image (OpenCV):', image_path)
+                os.remove(image_path)
+                continue
+
             tip = imghdr.what(image_path)
-            if tip not in image_exts: 
-                print('Image not in ext list {}'.format(image_path))
+            if tip not in image_exts:
+                print('Unsupported image type:', image_path)
                 os.remove(image_path)
+                continue
+
+            # Additional corruption check using PIL
+            try:
+                with Image.open(image_path) as im:
+                    im.verify()
+            except Exception as pil_err:
+                print('Corrupted image:', image_path)
+                os.remove(image_path)
+
         except Exception as e: 
-            print('Issue with image {}'.format(image_path))
-            # os.remove(image_path)
+            print('Error processing image:', image_path)
+            print(str(e))
+            os.remove(image_path)