Spaces:
Running
Running
import gradio as gr | |
import torch | |
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline | |
models=[ | |
"Nahrawy/AIorNot", | |
"RishiDarkDevil/ai-image-det-resnet152", | |
"arnolfokam/ai-generated-image-detector", | |
"umm-maybe/AI-image-detector", | |
] | |
#pipe = pipeline("image-classification", "umm-maybe/AI-image-detector") | |
def image_classifier(image): | |
outputs = pipe(image) | |
results = {} | |
for result in outputs: | |
results[result['label']] = result['score'] | |
return results | |
#demo = gr.Interface(fn=image_classifier, inputs=gr.Image(type="pil"), outputs="label", title=title, description=description) | |
#demo.launch(show_api=False) | |
def aiornot(image,mod_choose): | |
labels = ["Real", "AI"] | |
#feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50") | |
mod=models[int(mod_choose)] | |
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) | |
print(outputs) | |
print(dir(outputs)) | |
logits = outputs.logits | |
print (logits) | |
prediction = logits.argmax(-1).item() | |
print(prediction) | |
label = labels[prediction] | |
return label | |
with gr.Blocks() as app: | |
with gr.Row(): | |
with gr.Column(): | |
inp = gr.Image() | |
mod_choose=gr.Number(value=0) | |
btn = gr.Button() | |
outp = gr.Textbox() | |
btn.click(aiornot,[inp,mod_choose],outp) | |
app.launch() |