Update app.py

app.py

@@ -43,42 +43,62 @@ class AIDetector:
 
 def process_zip(zip_file):
     temp_dir = tempfile.mkdtemp()
-    with zipfile.ZipFile(zip_file.name, 'r') as z:
-        z.extractall(temp_dir)
     
-    labels, preds, images = [], [], []
-    false_positives, false_negatives = [], []
-    detector = AIDetector()
-    
-    for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
-        folder_path = os.path.join(temp_dir, folder_name)
-        if not os.path.exists(folder_path):
-            print(f"Folder not found: {folder_path}")
-            continue
-        for img_name in os.listdir(folder_path):
-            img_path = os.path.join(folder_path, img_name)
-            try:
-                img = Image.open(img_path).convert("RGB")
-                _, prediction = detector.predict(img)
-                pred_label = 0 if prediction["AI"] > prediction["Real"] else 1
-                
-                preds.append(pred_label)
-                labels.append(ground_truth_label)
-                images.append(img_name)
+    try:
+        # Validate zip structure
+        with zipfile.ZipFile(zip_file.name, 'r') as z:
+            file_list = z.namelist()
+            if not ('real/' in file_list and 'ai/' in file_list):
+                raise ValueError("Zip file must contain 'real' and 'ai' folders")
+            
+            z.extractall(temp_dir)
+        
+        labels, preds, images = [], [], []
+        false_positives, false_negatives = [], []
+        detector = AIDetector()
+        
+        total_images = sum(len(files) for _, _, files in os.walk(temp_dir))
+        processed_images = 0
+        
+        for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
+            folder_path = os.path.join(temp_dir, folder_name)
+            if not os.path.exists(folder_path):
+                raise ValueError(f"Folder not found: {folder_path}")
+            
+            for img_name in os.listdir(folder_path):
+                img_path = os.path.join(folder_path, img_name)
+                try:
+                    with Image.open(img_path).convert("RGB") as img:
+                        _, prediction = detector.predict(img)
+                    
+                    pred_label = 0 if prediction["AI"] > prediction["Real"] else 1
+                    
+                    preds.append(pred_label)
+                    labels.append(ground_truth_label)
+                    images.append(img_name)
 
-                # Collect false positives and false negatives with image data
-                if pred_label != ground_truth_label:
-                    img_data = base64.b64encode(open(img_path, "rb").read()).decode()
-                    if pred_label == 1 and ground_truth_label == 0:
-                        false_positives.append((img_name, img_data))
-                    elif pred_label == 0 and ground_truth_label == 1:
-                        false_negatives.append((img_name, img_data))
+                    # Collect false positives and false negatives with image data
+                    if pred_label != ground_truth_label:
+                        with open(img_path, "rb") as img_file:
+                            img_data = base64.b64encode(img_file.read()).decode()
+                        if pred_label == 1 and ground_truth_label == 0:
+                            false_positives.append((img_name, img_data))
+                        elif pred_label == 0 and ground_truth_label == 1:
+                            false_negatives.append((img_name, img_data))
 
-            except Exception as e:
-                print(f"Error processing image {img_name}: {e}")
+                except Exception as e:
+                    print(f"Error processing image {img_name}: {e}")
+                
+                processed_images += 1
+                gr.Progress(processed_images / total_images)
+        
+        return evaluate_model(labels, preds, false_positives, false_negatives)
+    
+    except Exception as e:
+        raise gr.Error(f"Error processing zip file: {str(e)}")
     
-    shutil.rmtree(temp_dir)
-    return evaluate_model(labels, preds, false_positives, false_negatives)
+    finally:
+        shutil.rmtree(temp_dir)
 
 def format_classification_report(labels, preds):
     # Convert the report string to a dictionary
@@ -260,7 +280,7 @@ detector = AIDetector()
 
 def create_gradio_interface():
     with gr.Blocks() as app:
-        gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)</h4></h1></center>""")
+        gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Demo)</h4></h1></center>""")
 
         with gr.Tabs():
             with gr.Tab("Single Image Detection"):
@@ -279,7 +299,12 @@ def create_gradio_interface():
                         output_label = gr.Label(label="Output")
 
             with gr.Tab("Batch Image Processing"):
-                zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
+                zip_file = gr.File(
+                    label="Upload Zip (Zip file must contain 2 folders, one labeled 'real' and one labled 'ai' )",
+                    file_types=[".zip"],
+                    file_count="single",
+                    max_file_size=1024  # Increased to 1024 MB (1 GB)
+                )
                 batch_btn = gr.Button("Process Batch")
 
                 with gr.Group():
@@ -300,7 +325,8 @@ def create_gradio_interface():
         batch_btn.click(
             process_zip,
             zip_file,
-            [output_acc, output_roc, output_report, output_plots, output_fp_fn]
+            [output_acc, output_roc, output_report, output_plots, output_fp_fn],
+            api_name="batch_process"
         )
 
     return app