Update app.py

app.py

@@ -10,6 +10,8 @@ from PIL import Image
 import tempfile
 import numpy as np
 import urllib.request
+import base64
+from io import BytesIO
 
 MODEL_NAME = "cmckinle/sdxl-flux-detector"
 LABELS = ["AI", "Real"]
@@ -45,6 +47,7 @@ def process_zip(zip_file):
         z.extractall(temp_dir)
     
     labels, preds, images = [], [], []
+    false_positives, false_negatives = [], []
     detector = AIDetector()
     
     for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
@@ -62,11 +65,20 @@ def process_zip(zip_file):
                 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))
+
             except Exception as e:
                 print(f"Error processing image {img_name}: {e}")
     
     shutil.rmtree(temp_dir)
-    return evaluate_model(labels, preds)
+    return evaluate_model(labels, preds, false_positives, false_negatives)
 
 def format_classification_report(labels, preds):
     # Convert the report string to a dictionary
@@ -166,7 +178,7 @@ def format_classification_report(labels, preds):
     
     return html
     
-def evaluate_model(labels, preds):
+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)
@@ -190,7 +202,49 @@ def evaluate_model(labels, preds):
     
     plt.tight_layout()
 
-    return accuracy, roc_score, report_html, fig
+    # Create HTML for false positives and negatives with images
+    fp_fn_html = """
+    <style>
+        .image-grid {
+            display: flex;
+            flex-wrap: wrap;
+            gap: 10px;
+        }
+        .image-item {
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+        }
+        .image-item img {
+            max-width: 200px;
+            max-height: 200px;
+        }
+    </style>
+    """
+
+    fp_fn_html += "<h3>False Positives (AI images classified as Real):</h3>"
+    fp_fn_html += '<div class="image-grid">'
+    for img_name, img_data in false_positives:
+        fp_fn_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>'
+
+    fp_fn_html += "<h3>False Negatives (Real images classified as AI):</h3>"
+    fp_fn_html += '<div class="image-grid">'
+    for img_name, img_data in false_negatives:
+        fp_fn_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>'
+
+    return accuracy, roc_score, report_html, fig, fp_fn_html
 
 def load_url(url):
     try:
@@ -234,6 +288,7 @@ def create_gradio_interface():
                     output_roc = gr.Label(label="ROC Score")
                     output_report = gr.HTML(label="Classification Report")
                     output_plots = gr.Plot(label="Confusion Matrix and ROC Curve")
+                    output_fp_fn = gr.HTML(label="False Positives and Negatives")
 
         load_btn.click(load_url, in_url, [inp, message])
         btn.click(
@@ -245,7 +300,7 @@ def create_gradio_interface():
         batch_btn.click(
             process_zip,
             zip_file,
-            [output_acc, output_roc, output_report, output_plots]
+            [output_acc, output_roc, output_report, output_plots, output_fp_fn]
         )
 
     return app