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)
        logits = outputs.logits
        probability = softmax(logits)
        px = pd.DataFrame(probability.numpy())
    prediction = logits.argmax(-1).item()
    label = labels[prediction]
    html_out = f"""
    <h1>This image is likely: {label}</h1><br><h3>
    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
    <br>    
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    return gr.HTML.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)
        logits = outputs.logits
        probability = softmax(logits)
        px = pd.DataFrame(probability.numpy())
    prediction = logits.argmax(-1).item()
    label = labels[prediction]
    html_out = f"""
    <h1>This image is likely: {label}</h1><br><h3>
    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
    <br>    
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    return gr.HTML.update(html_out)
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)
        logits = outputs.logits
        probability = softmax(logits)
        px = pd.DataFrame(probability.numpy())
    prediction = logits.argmax(-1).item()
    label = labels[prediction]
    html_out = f"""
    <h1>This image is likely: {label}</h1><br><h3>
    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
    <br>    
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    return gr.HTML.update(html_out)
with gr.Blocks() as app:
    with gr.Column():
        inp = gr.Image()
        btn = gr.Button()
    with gr.Group():        
        with gr.Row():
            with gr.Box():
                lab0 = gr.HTML(f"""Testing on Model: {models[0]}""")
                outp0 = gr.HTML("""""")
            with gr.Box():
                lab1 = gr.HTML(f"""Testing on Model: {models[1]}""")
                outp1 = gr.HTML("""""")
            with gr.Box():
                lab2 = gr.HTML(f"""Testing on Model: {models[2]}""")
                outp2 = gr.HTML("""""")            
    btn.click(aiornot0,[inp],outp0)
    btn.click(aiornot1,[inp],outp1)
    btn.click(aiornot2,[inp],outp2)
    
app.launch()