app.py

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
# 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):
            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")

                # Now use the specific model pipeline passed in
                pred = model_pipeline(img)
                pred_label = np.argmax([x['score'] for x in pred])

                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}")
    return labels, preds, images


# 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):
            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 = np.argmax([x['score'] for x in pred])

                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}")
    return labels, preds, images

# Batch processing for all models
def process_zip(zip_file):
    extracted_dir = extract_zip(zip_file.name)

    # Initialize model pipelines separately to avoid reuse issues
    model_pipelines = [pipe0, pipe1, pipe2]

    # Run classification for each model
    results = {}
    for idx, pipe in enumerate(model_pipelines):
        print(f"Processing with model {idx}")
        
        # Classify images with the correct pipeline per model
        labels, preds, images = classify_images(extracted_dir, pipe, idx)
        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 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 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)