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