Spaces:
Sleeping
Sleeping
| 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) | |