import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
from numpy import exp
import pandas as  pd

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_extractor0 = AutoFeatureExtractor.from_pretrained(mod)
    model0 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor0(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model0(**input)
        print (outputs)
        logits = outputs.logits
        print (logits)
        probability = softmax(logits)
        print(f'PROBABILITY ::: {probability}')
        print(probability[0][0])
        px = pd.DataFrame(probability.numpy())
        print(px)
        
    prediction = logits.argmax(-1).item()
    label = labels[prediction]
    
    html_out = f"""
    <h3>Model used: <a href src='https://huggingface.co/models/{mod}'>{mod}</a><br>
    This image is likely: {label}<br>
    Probabilites<br>
    AI: {px[0]}<br>
    Real: {px[1]}"""
    
    return gr.update(html_out)
def aiornot1(image):    
    labels = ["Real", "AI"]
    mod=models[1]
    feature_extractor1 = AutoFeatureExtractor.from_pretrained(mod)
    model1 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor1(image, return_tensors="pt")
    with torch.no_grad():
      outputs = model1(**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_extractor2 = AutoFeatureExtractor.from_pretrained(mod)
    model2 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor2(image, return_tensors="pt")
    with torch.no_grad():
      outputs = model2(**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(label=f'{models[0]}')
            outp0 = gr.HTML("""""")
            outp1 = gr.Textbox(label=f'{models[1]}')
            outp2 = gr.Textbox(label=f'{models[2]}')
    btn.click(aiornot0,[inp],outp0)
    btn.click(aiornot1,[inp],outp1)
    btn.click(aiornot2,[inp],outp2)
    
app.launch()