Update app.py

app.py

@@ -35,6 +35,115 @@ def softmax(vector):
     e = exp(vector - vector.max())  # for numerical stability
     return e / e.sum()
 
+# Single image classification functions
+def image_classifier0(image):
+    labels = ["AI", "Real"]
+    outputs = pipe0(image)
+    results = {}
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+    fin_sum.append(results)
+    return results
+
+def image_classifier1(image):
+    labels = ["AI", "Real"]
+    outputs = pipe1(image)
+    results = {}
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+    fin_sum.append(results)
+    return results
+
+def image_classifier2(image):
+    labels = ["AI", "Real"]
+    outputs = pipe2(image)
+    results = {}
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+    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
+
 # Function to extract images from zip
 def extract_zip(zip_file):
     temp_dir = tempfile.mkdtemp()  # Temporary directory
@@ -91,7 +200,7 @@ def evaluate_model(labels, preds):
 
     return accuracy, roc_score, report, fig, fig_roc
 
-# Gradio function for batch image processing with all models
+# Batch processing for all models
 def process_zip(zip_file):
     extracted_dir = extract_zip(zip_file.name)
     
@@ -118,34 +227,7 @@ def process_zip(zip_file):
             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):
-    labels = ["AI", "Real"]
-    outputs = pipe0(image)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
-    fin_sum.append(results)
-    return results
-
-def image_classifier1(image):
-    labels = ["AI", "Real"]
-    outputs = pipe1(image)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
-    fin_sum.append(results)
-    return results
-
-def image_classifier2(image):
-    labels = ["AI", "Real"]
-    outputs = pipe2(image)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
-    fin_sum.append(results)
-    return results
-
+# Single image section
 def load_url(url):
     try:
         urllib.request.urlretrieve(f'{url}', f"{uid}tmp_im.png")
@@ -201,29 +283,29 @@ 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(image_classifier0, [inp], [n_out0]).then(
-                image_classifier1, [inp], [n_out1]).then(
-                image_classifier2, [inp], [n_out2]).then(
+            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)
 
         # Tab for batch processing
         with gr.Tab("Batch Image Processing"):
             zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
-            
-            # 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")
 
+            for i, model in enumerate(models):
+                with gr.Group():
+                    gr.Markdown(f"### Results for {model}")
+                    globals()[f'output_acc{i}'] = gr.Label(label=f"Model {i} Accuracy")
+                    globals()[f'output_roc{i}'] = gr.Label(label=f"Model {i} ROC Score")
+                    globals()[f'output_report{i}'] = gr.Textbox(label=f"Model {i} Classification Report", lines=10)
+                    globals()[f'output_cm{i}'] = gr.Plot(label=f"Model {i} Confusion Matrix")
+                    globals()[f'output_roc_plot{i}'] = gr.Plot(label=f"Model {i} ROC Curve")
+
     # Connect batch processing
     batch_btn.click(process_zip, zip_file, 
-                    [output_acc, output_roc, output_report, output_cm, output_roc_plot] * 3)  # For all 3 models
+                    [output_acc0, output_roc0, output_report0, output_cm0, output_roc_plot0,
+                     output_acc1, output_roc1, output_report1, output_cm1, output_roc_plot1,
+                     output_acc2, output_roc2, output_report2, output_cm2, output_roc_plot2])
 
 app.launch(show_api=False, max_threads=24)