Create app.py

app.py

@@ -0,0 +1,331 @@
+import gradio as gr
+import torch
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
+import os
+import zipfile
+import shutil
+import matplotlib.pyplot as plt
+from sklearn.metrics import accuracy_score, roc_auc_score, confusion_matrix, classification_report, roc_curve, auc
+from tqdm import tqdm
+from PIL import Image
+import uuid
+import tempfile
+import pandas as pd
+from numpy import exp
+import numpy as np
+from sklearn.metrics import ConfusionMatrixDisplay
+import urllib.request
+
+# Define models
+models = [
+    "umm-maybe/AI-image-detector",
+    "Organika/sdxl-detector",
+    "cmckinle/sdxl-flux-detector",
+]
+
+pipe0 = pipeline("image-classification", f"{models[0]}")
+pipe1 = pipeline("image-classification", f"{models[1]}")
+pipe2 = pipeline("image-classification", f"{models[2]}")
+
+fin_sum = []
+uid = uuid.uuid4()
+
+# Softmax function
+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
+    with zipfile.ZipFile(zip_file, 'r') as z:
+        z.extractall(temp_dir)
+    return temp_dir
+
+# Function to classify images in a folder
+def classify_images(image_dir, model_pipeline, model_idx):
+    images = []
+    labels = []
+    preds = []
+    for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
+        folder_path = os.path.join(image_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")
+                
+                # Ensure that each image is being processed by the correct model pipeline
+                pred = model_pipeline(img)
+                pred_label = 0 if pred[0]['label'] == 'AI' else 1  # Assuming 'AI' is label 0 and 'Real' is label 1
+
+                preds.append(pred_label)
+                labels.append(ground_truth_label)
+                images.append(img_name)
+            except Exception as e:
+                print(f"Error processing image {img_name} in model {model_idx}: {e}")
+    
+    print(f"Model {model_idx} processed {len(images)} images")
+    return labels, preds, images
+
+# Function to generate evaluation metrics
+def evaluate_model(labels, preds):
+    cm = confusion_matrix(labels, preds)
+    accuracy = accuracy_score(labels, preds)
+    roc_score = roc_auc_score(labels, preds)
+    report = classification_report(labels, preds)
+    fpr, tpr, _ = roc_curve(labels, preds)
+    roc_auc = auc(fpr, tpr)
+
+    fig, ax = plt.subplots()
+    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=["AI", "Real"])
+    disp.plot(cmap=plt.cm.Blues, ax=ax)
+    plt.close(fig)
+
+    fig_roc, ax_roc = plt.subplots()
+    ax_roc.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
+    ax_roc.plot([0, 1], [0, 1], color='gray', linestyle='--')
+    ax_roc.set_xlim([0.0, 1.0])
+    ax_roc.set_ylim([0.0, 1.05])
+    ax_roc.set_xlabel('False Positive Rate')
+    ax_roc.set_ylabel('True Positive Rate')
+    ax_roc.set_title('Receiver Operating Characteristic (ROC) Curve')
+    ax_roc.legend(loc="lower right")
+    plt.close(fig_roc)
+
+    return accuracy, roc_score, report, fig, fig_roc
+
+# Batch processing for all models
+def process_zip(zip_file):
+    extracted_dir = extract_zip(zip_file.name)
+
+    # Run classification for each model
+    results = {}
+    for idx in range(len(models)):
+        print(f"Processing with model {models[idx]}")  # Debugging to show which model is being used
+
+        # Create a new pipeline for each model within the loop
+        pipe = pipeline("image-classification", f"{models[idx]}")  
+        print(f"Initialized pipeline for {models[idx]}")  # Confirm pipeline is initialized correctly
+
+        # Classify images with the correct pipeline per model
+        labels, preds, images = classify_images(extracted_dir, pipe, idx)
+        
+        # Debugging: Print the predictions to ensure they're different
+        print(f"Predictions for model {models[idx]}: {preds}")
+
+        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 results for all 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 section
+def load_url(url):
+    try:
+        urllib.request.urlretrieve(f'{url}', f"{uid}tmp_im.png")
+        image = Image.open(f"{uid}tmp_im.png")
+        mes = "Image Loaded"
+    except Exception as e:
+        image = None
+        mes = f"Image not Found<br>Error: {e}"
+    return image, mes
+
+def tot_prob():
+    try:
+        fin_out = sum([result["Real"] for result in fin_sum]) / len(fin_sum)
+        fin_sub = 1 - fin_out
+        out = {
+            "Real": f"{fin_out:.4f}",
+            "AI": f"{fin_sub:.4f}"
+        }
+        return out
+    except Exception as e:
+        print(e)
+        return None
+
+def fin_clear():
+    fin_sum.clear()
+    return None
+
+# 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>""")
+
+    with gr.Tabs():
+        # Tab for single image detection
+        with gr.Tab("Single Image Detection"):
+            with gr.Column():
+                inp = gr.Image(type='pil')
+                in_url = gr.Textbox(label="Image URL")
+                with gr.Row():
+                    load_btn = gr.Button("Load URL")
+                    btn = gr.Button("Detect AI")
+                mes = gr.HTML("""""")
+
+            with gr.Group():
+                with gr.Row():
+                    fin = gr.Label(label="Final Probability")
+                with gr.Row():
+                    for i, model in enumerate(models):
+                        with gr.Box():
+                            gr.HTML(f"""<b>Testing on Model {i}: <a href='https://huggingface.co/{model}'>{model}</a></b>""")
+                            globals()[f'outp{i}'] = gr.HTML("""""")
+                            globals()[f'n_out{i}'] = gr.Label(label="Output")
+
+            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)
+
+        # Tab for batch processing
+        with gr.Tab("Batch Image Processing"):
+            zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
+            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_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)