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) logits = outputs.logits print (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() mod_choose=gr.Number(value=0) btn = gr.Button() outp = gr.Textbox() btn.click(aiornot,[inp,mod_choose],outp) app.launch()