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"""

This image is likely: {label}


Probabilities:
Real: {float(px[1][0]):.4f}
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"""

This image is likely: {label}


Probabilities:
Real: {float(px[1][0]):.4f}
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"""

This image is likely: {label}


Probabilities:
Real: {float(px[1][0]):.4f}
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
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("""

AI Image Detector

(Test Demo - accuracy varies by model)

""") 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"""Testing on Model {i}: {model}""") 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)