| import gradio as gr | |
| import torch | |
| from transformers import AutoFeatureExtractor, AutoModelForImageClassification | |
| def aiornot(img): | |
| labels = ["Real", "AI"] | |
| 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(): | |
| inp = gr.Image() | |
| btn = gr.Button() | |
| outp = gr.Textbox() | |
| btn.click(aiornot,inp,outp) | |
| app.launch() |