Update app.py

app.py

@@ -91,13 +91,32 @@ def evaluate_model(labels, preds):
 
     return accuracy, roc_score, report, fig, fig_roc
 
-# Gradio function for batch image processing
+# Gradio function for batch image processing with all models
 def process_zip(zip_file):
     extracted_dir = extract_zip(zip_file.name)
-    labels, preds, images = classify_images(extracted_dir, pipe0)  # You can switch to pipe1 or pipe2
-    accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds)
+    
+    # Run classification for each model
+    results = {}
+    for idx, pipe in enumerate([pipe0, pipe1, pipe2]):
+        labels, preds, images = classify_images(extracted_dir, pipe)
+        accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds)
+        
+        # Store results for each model
+        results[f'Model_{idx}_accuracy'] = accuracy
+        results[f'Model_{idx}_roc_score'] = roc_score
+        results[f'Model_{idx}_report'] = report
+        results[f'Model_{idx}_cm_fig'] = cm_fig
+        results[f'Model_{idx}_roc_fig'] = roc_fig
+
     shutil.rmtree(extracted_dir)  # Clean up extracted files
-    return accuracy, roc_score, report, cm_fig, roc_fig
+
+    # Return results for all three models
+    return (results['Model_0_accuracy'], results['Model_0_roc_score'], results['Model_0_report'], 
+            results['Model_0_cm_fig'], results['Model_0_roc_fig'],
+            results['Model_1_accuracy'], results['Model_1_roc_score'], results['Model_1_report'], 
+            results['Model_1_cm_fig'], results['Model_1_roc_fig'],
+            results['Model_2_accuracy'], results['Model_2_roc_score'], results['Model_2_report'], 
+            results['Model_2_cm_fig'], results['Model_2_roc_fig'])
 
 # Single image classification functions
 def image_classifier0(image):
@@ -127,87 +146,6 @@ def image_classifier2(image):
     fin_sum.append(results)
     return results
 
-def aiornot0(image):
-    labels = ["AI", "Real"]
-    mod = models[0]
-    feature_extractor0 = AutoFeatureExtractor.from_pretrained(mod)
-    model0 = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor0(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model0(**input)
-        logits = outputs.logits
-        probability = softmax(logits)  # Apply softmax on logits
-        px = pd.DataFrame(probability.numpy())
-    prediction = logits.argmax(-1).item()
-    label = labels[prediction]
-
-    html_out = f"""
-    <h1>This image is likely: {label}</h1><br><h3>
-    Probabilities:<br>
-    Real: {float(px[1][0]):.4f}<br>
-    AI: {float(px[0][0]):.4f}"""
-    
-    results = {
-        "Real": float(px[1][0]),
-        "AI": float(px[0][0])
-    }
-    fin_sum.append(results)
-    return gr.HTML.update(html_out), results
-
-def aiornot1(image):
-    labels = ["AI", "Real"]
-    mod = models[1]
-    feature_extractor1 = AutoFeatureExtractor.from_pretrained(mod)
-    model1 = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor1(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model1(**input)
-        logits = outputs.logits
-        probability = softmax(logits)  # Apply softmax on logits
-        px = pd.DataFrame(probability.numpy())
-    prediction = logits.argmax(-1).item()
-    label = labels[prediction]
-
-    html_out = f"""
-    <h1>This image is likely: {label}</h1><br><h3>
-    Probabilities:<br>
-    Real: {float(px[1][0]):.4f}<br>
-    AI: {float(px[0][0]):.4f}"""
-    
-    results = {
-        "Real": float(px[1][0]),
-        "AI": float(px[0][0])
-    }
-    fin_sum.append(results)
-    return gr.HTML.update(html_out), results
-
-def aiornot2(image):
-    labels = ["AI", "Real"]
-    mod = models[2]
-    feature_extractor2 = AutoFeatureExtractor.from_pretrained(mod)
-    model2 = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor2(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model2(**input)
-        logits = outputs.logits
-        probability = softmax(logits)  # Apply softmax on logits
-        px = pd.DataFrame(probability.numpy())
-    prediction = logits.argmax(-1).item()
-    label = labels[prediction]
-
-    html_out = f"""
-    <h1>This image is likely: {label}</h1><br><h3>
-    Probabilities:<br>
-    Real: {float(px[1][0]):.4f}<br>
-    AI: {float(px[0][0]):.4f}"""
-    
-    results = {
-        "Real": float(px[1][0]),
-        "AI": float(px[0][0])
-    }
-    fin_sum.append(results)
-    return gr.HTML.update(html_out), results
-
 def load_url(url):
     try:
         urllib.request.urlretrieve(f'{url}', f"{uid}tmp_im.png")
@@ -235,12 +173,6 @@ def fin_clear():
     fin_sum.clear()
     return None
 
-def upd(image):
-    rand_im = uuid.uuid4()
-    image.save(f"{rand_im}-vid_tmp_proc.png")
-    out = Image.open(f"{rand_im}-vid_tmp_proc.png")
-    return out
-
 # Set up Gradio app
 with gr.Blocks() as app:
     gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)</h4></h1></center>""")
@@ -269,11 +201,6 @@ with gr.Blocks() as app:
             btn.click(fin_clear, None, fin, show_progress=False)
             load_btn.click(load_url, in_url, [inp, mes])
 
-            btn.click(aiornot0, [inp], [outp0, n_out0]).then(
-                aiornot1, [inp], [outp1, n_out1]).then(
-                aiornot2, [inp], [outp2, n_out2]).then(
-                tot_prob, None, fin, show_progress=False)
-
             btn.click(image_classifier0, [inp], [n_out0]).then(
                 image_classifier1, [inp], [n_out1]).then(
                 image_classifier2, [inp], [n_out2]).then(
@@ -282,15 +209,21 @@ with gr.Blocks() as app:
         # Tab for batch processing
         with gr.Tab("Batch Image Processing"):
             zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
-            output_acc = gr.Label(label="Accuracy")
-            output_roc = gr.Label(label="ROC Score")
-            output_report = gr.Textbox(label="Classification Report", lines=10)
-            output_cm = gr.Plot(label="Confusion Matrix")
-            output_roc_plot = gr.Plot(label="ROC Curve")
+            
+            # Outputs for all three models
+            for i in range(3):
+                with gr.Group():
+                    gr.Markdown(f"### Results for Model {i}")
+                    output_acc = gr.Label(label=f"Model {i} Accuracy")
+                    output_roc = gr.Label(label=f"Model {i} ROC Score")
+                    output_report = gr.Textbox(label=f"Model {i} Classification Report", lines=10)
+                    output_cm = gr.Plot(label=f"Model {i} Confusion Matrix")
+                    output_roc_plot = gr.Plot(label=f"Model {i} ROC Curve")
 
             batch_btn = gr.Button("Process Batch")
 
     # Connect batch processing
-    batch_btn.click(process_zip, zip_file, [output_acc, output_roc, output_report, output_cm, output_roc_plot])
+    batch_btn.click(process_zip, zip_file, 
+                    [output_acc, output_roc, output_report, output_cm, output_roc_plot] * 3)  # For all 3 models
 
-app.launch(show_api=False, max_threads=24)
+app.launch(show_api=False, max_threads=24)