import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
from numpy import exp
 
def softmax(vector):
 e = exp(vector)
 return e / e.sum()

    
models=[
    "Nahrawy/AIorNot",
    "arnolfokam/ai-generated-image-detector",
    "umm-maybe/AI-image-detector",
]

def aiornot0(image):    
    labels = ["Real", "AI"]
    mod=models[0]
    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)
        print (outputs)
        logits = outputs.logits
        print (logits)
        probability = softmax(logits, axis=-1)
        print(probability)

    prediction = logits.argmax(-1).item()
    label = labels[prediction] 
    return label
def aiornot1(image):    
    labels = ["Real", "AI"]
    mod=models[1]
    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)
      print (outputs)
      logits = outputs.logits
      print (logits)
    prediction = logits.argmax(-1).item()
    label = labels[prediction] 
    return label
def aiornot2(image):    
    labels = ["Real", "AI"]
    mod=models[2]
    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)
      print (outputs)
      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()
            
        with gr.Column():
            outp0 = gr.Textbox()
            outp1 = gr.Textbox()
            outp2 = gr.Textbox()
    btn.click(aiornot0,[inp],outp0)
    btn.click(aiornot1,[inp],outp1)
    btn.click(aiornot2,[inp],outp2)
    
app.launch()