|
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", |
|
] |
|
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 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def aiornot(image,mod_choose): |
|
labels = ["Real", "AI"] |
|
|
|
|
|
mod=models[int(mod_choose)] |
|
feature_extractor = AutoFeatureExtractor.from_pretrained("Nahrawy/AIorNot") |
|
model = AutoModelForImageClassification.from_pretrained("Nahrawy/AIorNot") |
|
|
|
input = feature_extractor(image, return_tensors="pt") |
|
with torch.no_grad(): |
|
outputs = model(**input) |
|
logits = outputs.logits |
|
prediction = logits.argmax(-1).item() |
|
label = labels[prediction] |
|
return label |
|
|
|
with gr.Blocks() as app: |
|
with gr.Row(): |
|
with gr.Column(): |
|
inp = gr.Image(type='filepath') |
|
mod_choose=gr.Number(value=0) |
|
btn = gr.Button() |
|
outp = gr.Textbox() |
|
btn.click(aiornot,[inp,mod_choose],outp) |
|
app.launch() |