import gradio as gr import torch from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline #from transformers import pipeline import os from numpy import exp import pandas as pd from PIL import Image import urllib.request import uuid uid=uuid.uuid4() models=[ "Nahrawy/AIorNot", "umm-maybe/AI-image-detector", "Nahrawy/AIorNot", #"arnolfokam/ai-generated-image-detector", ] pipe0 = pipeline("image-classification", f"{models[0]}") pipe1 = pipeline("image-classification", f"{models[1]}") pipe2 = pipeline("image-classification", f"{models[2]}") fin_sum=[] def image_classifier0(image): labels = ["AI","Real"] outputs = pipe0(image) results = {} result_test={} for idx,result in enumerate(outputs): results[labels[idx]] = outputs[idx]['score'] #print (result_test) #for result in outputs: # results[result['label']] = result['score'] #print (results) fin_sum.append(results) return results def image_classifier1(image): labels = ["AI","Real"] outputs = pipe1(image) results = {} result_test={} for idx,result in enumerate(outputs): results[labels[idx]] = outputs[idx]['score'] #print (result_test) #for result in outputs: # results[result['label']] = result['score'] #print (results) fin_sum.append(results) return results def image_classifier2(image): labels = ["AI","Real"] outputs = pipe2(image) results = {} result_test={} for idx,result in enumerate(outputs): results[labels[idx]] = outputs[idx]['score'] #print (result_test) #for result in outputs: # results[result['label']] = result['score'] #print (results) fin_sum.append(results) return results def softmax(vector): e = exp(vector) return e / e.sum() 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) px = pd.DataFrame(probability.numpy()) prediction = logits.argmax(-1).item() label = labels[prediction] html_out = f"""

This image is likely: {label}


Probabilites:
Real: {px[1][0]}
AI: {px[0][0]}""" results = {} for idx,result in enumerate(px): results[labels[idx]] = px[idx][0] #results[labels['label']] = result['score'] 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) px = pd.DataFrame(probability.numpy()) prediction = logits.argmax(-1).item() label = labels[prediction] html_out = f"""

This image is likely: {label}


Probabilites:
Real: {px[1][0]}
AI: {px[0][0]}""" results = {} for idx,result in enumerate(px): results[labels[idx]] = px[idx][0] #results[labels['label']] = result['score'] fin_sum.append(results) return gr.HTML.update(html_out),results def aiornot2(image): labels = ["Real", "AI"] mod=models[2] feature_extractor2 = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50") 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) px = pd.DataFrame(probability.numpy()) prediction = logits.argmax(-1).item() label = labels[prediction] html_out = f"""

This image is likely: {label}


Probabilites:
Real: {px[0][0]}
AI: {px[1][0]}""" results = {} for idx,result in enumerate(px): results[labels[idx]] = px[idx][0] #results[labels['label']] = result['score'] fin_sum.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
Error: {e}" return image,mes def tot_prob(): try: fin_out = fin_sum[0]["Real"]+fin_sum[1]["Real"]+fin_sum[2]["Real"]+fin_sum[3]["Real"]+fin_sum[4]["Real"]+fin_sum[5]["Real"] fin_out = fin_out/6 fin_sub = 1-fin_out out={ "Real":f"{fin_out}", "AI":f"{fin_sub}" } #fin_sum.clear() #print (fin_out) return out except Exception as e: pass print (e) return None def fin_clear(): fin_sum.clear() return None def upd(image): print (image) rand_im = uuid.uuid4() image.save(f"{rand_im}-vid_tmp_proc.png") out = Image.open(f"{rand_im}-vid_tmp_proc.png") #image.save(f"{rand_im}-vid_tmp_proc.png") #out = os.path.abspath(f"{rand_im}-vid_tmp_proc.png") #out_url = f'https://omnibus_AI_or_Not_dev.hf.space/file={out}' #out_url = f"{rand_im}-vid_tmp_proc.png" return out with gr.Blocks() as app: gr.Markdown("""

AI Image Detector

(Test Demo - accuracy varies by model)""") 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(): lab0 = gr.HTML(f"""Testing on Model: {models[0]}""") nun0 = gr.HTML("""""") with gr.Box(): lab1 = gr.HTML(f"""Testing on Model: {models[1]}""") nun1 = gr.HTML("""""") with gr.Box(): lab2 = gr.HTML(f"""Testing on Model: {models[2]}""") nun2 = gr.HTML("""""") with gr.Row(): with gr.Box(): n_out0=gr.Label(label="Output") outp0 = gr.HTML("""""") with gr.Box(): n_out1=gr.Label(label="Output") outp1 = gr.HTML("""""") with gr.Box(): n_out2=gr.Label(label="Output") outp2 = gr.HTML("""""") with gr.Row(): with gr.Box(): n_out3=gr.Label(label="Output") outp3 = gr.HTML("""""") with gr.Box(): n_out4=gr.Label(label="Output") outp4 = gr.HTML("""""") with gr.Box(): n_out5=gr.Label(label="Output") outp5 = gr.HTML("""""") hid_box=gr.Textbox(visible=False) hid_im = gr.Image(type="pil",visible=False) def echo(inp): return inp #inp.change(echo,inp,hid_im).then(upd,hid_im,inp) 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(tot_prob,None,fin,show_progress=False) btn.click(aiornot1,[inp],[outp1,n_out1]).then(tot_prob,None,fin,show_progress=False) btn.click(aiornot2,[inp],[outp2,n_out2]).then(tot_prob,None,fin,show_progress=False) btn.click(image_classifier0,[inp],[n_out3]).then(tot_prob,None,fin,show_progress=False) btn.click(image_classifier1,[inp],[n_out4]).then(tot_prob,None,fin,show_progress=False) btn.click(image_classifier2,[inp],[n_out5]).then(tot_prob,None,fin,show_progress=False) app.launch(show_api=False,max_threads=24)