Update app.py

app.py

@@ -10,12 +10,17 @@ import zipfile
 dataset_dir = "dataset"
 class_root_dir = None 
 def find_class_root(base_dir):
-    # Look for the first folder that contains subfolders (i.e., class folders)
     for root, dirs, files in os.walk(base_dir):
-        if all(os.path.isdir(os.path.join(root, d)) for d in dirs) and len(dirs) > 0:
-            return root
+        # Check if this folder contains class folders (i.e., folders with image files)
+        if any(os.path.isdir(os.path.join(root, d)) for d in dirs):
+            # Check if those subfolders contain image files
+            for d in dirs:
+                subfolder = os.path.join(root, d)
+                if any(fname.lower().endswith(('.png', '.jpg', '.jpeg')) for fname in os.listdir(subfolder)):
+                    return root
     return None
 
+
 def upload_dataset(file):
     global class_root_dir
     os.makedirs(dataset_dir, exist_ok=True)