AI-or-Not-v2 / app.py
Omnibus's picture
Update app.py
1573304
raw
history blame
No virus
1.09 kB
import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
models=[
"Nahrawy/AIorNot",
"RishiDarkDevil/ai-image-det-resnet152",
"arnolfokam/ai-generated-image-detector",
]
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("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()