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()