|
|
import gradio as gr |
|
|
import torch |
|
|
from transformers import AutoFeatureExtractor, AutoModelForImageClassification |
|
|
from numpy import exp |
|
|
import pandas as pd |
|
|
from PIL import Image |
|
|
import urllib.request |
|
|
import uuid |
|
|
|
|
|
uid = uuid.uuid4() |
|
|
models = [ |
|
|
"cmckinle/sdxl-flux-detector", |
|
|
"umm-maybe/AI-image-detector", |
|
|
"Organika/sdxl-detector", |
|
|
] |
|
|
results_store = [] |
|
|
|
|
|
def softmax(vector): |
|
|
e = exp(vector) |
|
|
return e / e.sum() |
|
|
|
|
|
def aiornot(image, model_index): |
|
|
mod = models[model_index] |
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained(mod) |
|
|
model = AutoModelForImageClassification.from_pretrained(mod) |
|
|
input = feature_extractor(image, return_tensors="pt") |
|
|
with torch.no_grad(): |
|
|
outputs = model(**input) |
|
|
logits = outputs.logits |
|
|
probability = softmax(logits) |
|
|
px = pd.DataFrame(probability.numpy()) |
|
|
|
|
|
if model_index == 2: |
|
|
real_prob, ai_prob = px[0][0], px[1][0] |
|
|
label = "Real" if real_prob > ai_prob else "AI" |
|
|
else: |
|
|
ai_prob, real_prob = px[0][0], px[1][0] |
|
|
label = "AI" if ai_prob > real_prob else "Real" |
|
|
|
|
|
html_out = f""" |
|
|
<h1>This image is likely: {label}</h1><br><h3> |
|
|
Probabilities:<br> |
|
|
Real: {real_prob:.4f}<br> |
|
|
AI: {ai_prob:.4f}""" |
|
|
|
|
|
results = {"Real": real_prob, "AI": ai_prob} |
|
|
results_store.append(results) |
|
|
return gr.HTML.update(html_out), results |
|
|
|
|
|
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 calculate_final_prob(): |
|
|
if not results_store: |
|
|
return {"Real": "N/A", "AI": "N/A"} |
|
|
fin_out = sum(result["Real"] for result in results_store) / len(results_store) |
|
|
return { |
|
|
"Real": f"{fin_out:.4f}", |
|
|
"AI": f"{1 - fin_out:.4f}" |
|
|
} |
|
|
|
|
|
with gr.Blocks() as app: |
|
|
gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)</h4></center>""") |
|
|
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.Column(): |
|
|
gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{model}'>{model}</a></b>""") |
|
|
output_html = gr.HTML() |
|
|
output_label = gr.Label(label="Output") |
|
|
btn.click(aiornot, inputs=[inp, gr.Number(value=i, visible=False)], outputs=[output_html, output_label]) |
|
|
if i == len(models) - 1: |
|
|
btn.click(lambda: results_store.clear(), outputs=None) |
|
|
btn.click(calculate_final_prob, outputs=fin) |
|
|
load_btn.click(load_url, in_url, [inp, mes]) |
|
|
|
|
|
app.launch(show_api=False, max_threads=24) |
