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 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 model model = "cmckinle/sdxl-flux-detector" pipe = pipeline("image-classification", model) 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 function def image_classifier(image): labels = ["AI", "Real"] outputs = pipe(image) results = {} for idx, result in enumerate(outputs): results[labels[idx]] = float(outputs[idx]['score']) fin_sum.append(results) return results def aiornot(image): labels = ["AI", "Real"] feature_extractor = AutoFeatureExtractor.from_pretrained(model) model_cls = AutoModelForImageClassification.from_pretrained(model) input = feature_extractor(image, return_tensors="pt") with torch.no_grad(): outputs = model_cls(**input) logits = outputs.logits probability = softmax(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() 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): 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") pred = pipe(img) pred_label = 0 if pred[0]['label'] == 'AI' else 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}: {e}") print(f"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 def process_zip(zip_file): extracted_dir = extract_zip(zip_file.name) labels, preds, images = classify_images(extracted_dir) accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds) shutil.rmtree(extracted_dir) # Clean up extracted files return accuracy, roc_score, report, cm_fig, 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(): with gr.Box(): gr.HTML(f"""Testing on Model: {model}""") outp = gr.HTML("""""") n_out = gr.Label(label="Output") btn.click(fin_clear, None, fin, show_progress=False) load_btn.click(load_url, in_url, [inp, mes]) btn.click(aiornot, [inp], [outp, n_out]).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") with gr.Group(): gr.Markdown(f"### Results for {model}") 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") # Connect batch processing batch_btn.click(process_zip, zip_file, [output_acc, output_roc, output_report, output_cm, output_roc_plot]) app.launch(show_api=False, max_threads=24)