Update app.py

app.py

@@ -53,46 +53,7 @@ def process_zip(zip_file):
             
             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:
-                        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}")
-                
-                processed_images += 1
-                gr.Progress(processed_images / total_images)
-        
-        return evaluate_model(labels, preds, false_positives, false_negatives)
+        return evaluate_model(temp_dir)
     
     except Exception as e:
         raise gr.Error(f"Error processing zip file: {str(e)}")
@@ -100,11 +61,100 @@ def process_zip(zip_file):
     finally:
         shutil.rmtree(temp_dir)
 
+def process_files(ai_files, real_files):
+    temp_dir = tempfile.mkdtemp()
+    try:
+        # Process AI files
+        ai_folder = os.path.join(temp_dir, 'ai')
+        os.makedirs(ai_folder)
+        for file in ai_files:
+            shutil.copy(file.name, os.path.join(ai_folder, os.path.basename(file.name)))
+
+        # Process Real files
+        real_folder = os.path.join(temp_dir, 'real')
+        os.makedirs(real_folder)
+        for file in real_files:
+            shutil.copy(file.name, os.path.join(real_folder, os.path.basename(file.name)))
+
+        return evaluate_model(temp_dir)
+    except Exception as e:
+        raise gr.Error(f"Error processing individual files: {str(e)}")
+    finally:
+        shutil.rmtree(temp_dir)
+
+def evaluate_model(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:
+                    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}")
+            
+            processed_images += 1
+            gr.Progress(processed_images / total_images)
+    
+    return calculate_metrics(labels, preds, false_positives, false_negatives)
+
+def calculate_metrics(labels, preds, false_positives, false_negatives):
+    cm = confusion_matrix(labels, preds)
+    accuracy = accuracy_score(labels, preds)
+    roc_score = roc_auc_score(labels, preds)
+    report_html = format_classification_report(labels, preds)
+    fpr, tpr, _ = roc_curve(labels, preds)
+    roc_auc = auc(fpr, tpr)
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
+    
+    ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=LABELS).plot(cmap=plt.cm.Blues, ax=ax1)
+    ax1.set_title("Confusion Matrix")
+    
+    ax2.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
+    ax2.plot([0, 1], [0, 1], color='gray', linestyle='--')
+    ax2.set_xlim([0.0, 1.0])
+    ax2.set_ylim([0.0, 1.05])
+    ax2.set_xlabel('False Positive Rate')
+    ax2.set_ylabel('True Positive Rate')
+    ax2.set_title('ROC Curve')
+    ax2.legend(loc="lower right")
+    
+    plt.tight_layout()
+
+    fp_fn_html = create_fp_fn_html(false_positives, false_negatives)
+
+    return accuracy, roc_score, report_html, fig, fp_fn_html
+
 def format_classification_report(labels, preds):
-    # Convert the report string to a dictionary
     report_dict = classification_report(labels, preds, output_dict=True)
     
-    # Create an HTML table with updated CSS
     html = """
     <style>
         .report-table {
@@ -170,7 +220,6 @@ def format_classification_report(labels, preds):
         </tr>
     """
     
-    # Add rows for each class
     for class_name in ['0', '1']:
         html += f"""
         <tr>
@@ -182,7 +231,6 @@ def format_classification_report(labels, preds):
         </tr>
         """
     
-    # Add summary rows
     html += f"""
         <tr>
             <td>Accuracy</td>
@@ -207,33 +255,9 @@ def format_classification_report(labels, preds):
     """
     
     return html
-    
-def evaluate_model(labels, preds, false_positives, false_negatives):
-    cm = confusion_matrix(labels, preds)
-    accuracy = accuracy_score(labels, preds)
-    roc_score = roc_auc_score(labels, preds)
-    report_html = format_classification_report(labels, preds)
-    fpr, tpr, _ = roc_curve(labels, preds)
-    roc_auc = auc(fpr, tpr)
-
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
-    
-    ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=LABELS).plot(cmap=plt.cm.Blues, ax=ax1)
-    ax1.set_title("Confusion Matrix")
-    
-    ax2.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
-    ax2.plot([0, 1], [0, 1], color='gray', linestyle='--')
-    ax2.set_xlim([0.0, 1.0])
-    ax2.set_ylim([0.0, 1.05])
-    ax2.set_xlabel('False Positive Rate')
-    ax2.set_ylabel('True Positive Rate')
-    ax2.set_title('ROC Curve')
-    ax2.legend(loc="lower right")
-    
-    plt.tight_layout()
 
-    # Create HTML for false positives and negatives with images
-    fp_fn_html = """
+def create_fp_fn_html(false_positives, false_negatives):
+    html = """
     <style>
         .image-grid {
             display: flex;
@@ -252,29 +276,29 @@ def evaluate_model(labels, preds, false_positives, false_negatives):
     </style>
     """
 
-    fp_fn_html += "<h3>False Positives (AI images classified as Real):</h3>"
-    fp_fn_html += '<div class="image-grid">'
+    html += "<h3>False Positives (AI images classified as Real):</h3>"
+    html += '<div class="image-grid">'
     for img_name, img_data in false_positives:
-        fp_fn_html += f'''
+        html += f'''
         <div class="image-item">
             <img src="data:image/jpeg;base64,{img_data}" alt="{img_name}">
             <p>{img_name}</p>
         </div>
         '''
-    fp_fn_html += '</div>'
+    html += '</div>'
 
-    fp_fn_html += "<h3>False Negatives (Real images classified as AI):</h3>"
-    fp_fn_html += '<div class="image-grid">'
+    html += "<h3>False Negatives (Real images classified as AI):</h3>"
+    html += '<div class="image-grid">'
     for img_name, img_data in false_negatives:
-        fp_fn_html += f'''
+        html += f'''
         <div class="image-item">
             <img src="data:image/jpeg;base64,{img_data}" alt="{img_name}">
             <p>{img_name}</p>
         </div>
         '''
-    fp_fn_html += '</div>'
+    html += '</div>'
 
-    return accuracy, roc_score, report_html, fig, fp_fn_html
+    return html
 
 def load_url(url):
     try:
@@ -309,13 +333,28 @@ def create_gradio_interface():
                         output_label = gr.Label(label="Output")
 
             with gr.Tab("Batch Image Processing"):
-                zip_file = gr.File(
-                    label="Upload Zip (must contain 'real' and 'ai' folders)",
-                    file_types=[".zip"],
-                    file_count="single",
-                    max_file_size=1024  # 1024 MB (1 GB)
-                )
-                batch_btn = gr.Button("Process Batch", interactive=False)
+                with gr.Accordion("Upload Zip File (max 100MB)", open=False):
+                    zip_file = gr.File(
+                        label="Upload Zip (must contain 'real' and 'ai' folders)",
+                        file_types=[".zip"],
+                        file_count="single",
+                        max_file_size=100  # 100 MB limit
+                    )
+                    zip_process_btn = gr.Button("Process Zip", interactive=False)
+
+                with gr.Accordion("Upload Individual Files (for datasets over 100MB)", open=False):
+                    with gr.Row():
+                        ai_files = gr.File(
+                            label="Upload AI Images",
+                            file_types=["image"],
+                            file_count="multiple"
+                        )
+                        real_files = gr.File(
+                            label="Upload Real Images",
+                            file_types=["image"],
+                            file_count="multiple"
+                        )
+                    individual_process_btn = gr.Button("Process Individual Files", interactive=False)
 
                 with gr.Group():
                     gr.Markdown(f"### Results for {MODEL_NAME}")
@@ -332,28 +371,7 @@ def create_gradio_interface():
             [output_html, output_label]
         )
 
-        def enable_batch_btn(file):
+        def enable_zip_btn(file):
             return gr.Button.update(interactive=file is not None)
 
-        zip_file.upload(
-            enable_batch_btn,
-            zip_file,
-            batch_btn
-        )
-
-        batch_btn.click(
-            process_zip,
-            zip_file,
-            [output_acc, output_roc, output_report, output_plots, output_fp_fn],
-            api_name="batch_process"
-        )
-
-    return app
-
-if __name__ == "__main__":
-    app = create_gradio_interface()
-    app.launch(
-    show_api=False,
-    max_threads=24,
-    show_error=True
-)
+        def enable_individual_btn