import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
from numpy import exp
import pandas as  pd
from PIL import Image
import urllib.request 
import uuid
uid=uuid.uuid4()



def softmax(vector):
 e = exp(vector)
 return e / e.sum()

    
models=[
    "Nahrawy/AIorNot",
    "umm-maybe/AI-image-detector",
    "arnolfokam/ai-generated-image-detector",
    "Binyamin/Hybrid_1",
    "HuggingSara/model_soups",
    "psyne/AIResnetClone",

]

fin_sum=[]
#fin_res={f'{uid}':''}
#fin_sum.append(fin_res)
#tmp_res=
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>
 
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    results = {}
    for idx,result in enumerate(px):
        results[labels[idx]] = px[idx][0]
    #results[labels['label']] = result['score']
    #tmp_res={f'{uid}-0':results}
    fin_sum.append(results)
    return gr.HTML.update(html_out),results
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>
  
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    results = {}
    for idx,result in enumerate(px):
        results[labels[idx]] = px[idx][0]
    #results[labels['label']] = result['score']
    #tmp_res={f'{uid}-1':results}
    fin_sum.append(results)        
    return gr.HTML.update(html_out),results    
def aiornot2(image):    
    labels = ["AI", "Real"]
    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>
  
    Probabilites:<br>
    Real: {px[1][0]}<br>
    AI: {px[0][0]}"""

    results = {}
    for idx,result in enumerate(px):
        results[labels[idx]] = px[idx][0]
    #tmp_res={f'{uid}-2':results}
    fin_sum.append(results)
    return gr.HTML.update(html_out),results
def aiornot3(image):    
    labels = ["Real", "AI"]
    mod=models[3]
    feature_extractor3 = AutoFeatureExtractor.from_pretrained(mod)
    model3 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor3(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model3(**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>
 
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    results = {}
    for idx,result in enumerate(px):
        results[labels[idx]] = px[idx][0]
    #results[labels['label']] = result['score']
    return gr.HTML.update(html_out),results
def aiornot4(image):    
    labels = ["Real", "AI"]
    mod=models[4]
    feature_extractor4 = AutoFeatureExtractor.from_pretrained(mod)
    model4 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor4(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model4(**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>
  
    Probabilites:<br>
    Real: {px[0][0]}<br>
    AI: {px[1][0]}"""
    results = {}
    for idx,result in enumerate(px):
        results[labels[idx]] = px[idx][0]
    #results[labels['label']] = result['score']
    return gr.HTML.update(html_out),results    
def aiornot5(image):    
    labels = ["AI", "Real"]
    mod=models[5]
    feature_extractor5 = AutoFeatureExtractor.from_pretrained(mod)
    model5 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor5(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model5(**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>
  
    Probabilites:<br>
    Real: {px[1][0]}<br>
    AI: {px[0][0]}"""

    results = {}
    for idx,result in enumerate(px):
        results[labels[idx]] = px[idx][0]
    #results[labels['label']] = result['score']
    return gr.HTML.update(html_out),results    
def load_url(url):
    try:
        urllib.request.urlretrieve( 
            f'{url}', 
            f"{uid}tmp_im.png")         
        image = Image.open(f"{uid}tmp_im.png")
        mes = "Image Loaded"
    except Exception as e:
        image=None
        mes=f"Image not Found<br>Error: {e}"
    return image,mes

def tot_prob():
    try:
        fin_out = fin_sum[0]["Real"]+fin_sum[1]["Real"]+fin_sum[2]["Real"]
        print (fin_out)
    except Exception as e:
        print (f'ERROR :: {e}')
    
with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            in_url=gr.Textbox(label="Image URL")
            with gr.Row():
                load_btn=gr.Button("Load URL")
                btn = gr.Button("Detect AI")
            mes = gr.HTML("""""")
        inp = gr.Pil()
    with gr.Group():        
        with gr.Row():
            with gr.Box():
                lab0 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[0]}'>{models[0]}</a></b>""")
                n_out0=gr.Label(label="Output")
                outp0 = gr.HTML("""""")
            with gr.Box():
                lab1 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[1]}'>{models[1]}</a></b>""")
                n_out1=gr.Label(label="Output")
                outp1 = gr.HTML("""""")
            with gr.Box():
                lab2 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[2]}'>{models[2]}</a></b>""")
                n_out2=gr.Label(label="Output")
                outp2 = gr.HTML("""""")    
        with gr.Row():
            with gr.Box():
                lab3 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[3]}'>{models[3]}</a></b>""")
                n_out3=gr.Label(label="Output")
                outp3 = gr.HTML("""""")
            with gr.Box():
                lab4 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[4]}'>{models[4]}</a></b>""")
                n_out4=gr.Label(label="Output")
                outp4 = gr.HTML("""""")
            with gr.Box():
                lab5 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[5]}'>{models[5]}</a></b>""")
                n_out5=gr.Label(label="Output")
                outp5 = gr.HTML("""""")                    
    load_btn.click(load_url,in_url,[inp,mes])
    btn.click(aiornot0,[inp],[outp0,n_out0]).then(tot_prob,None,None)
    btn.click(aiornot1,[inp],[outp1,n_out1]).then(tot_prob,None,None)
    btn.click(aiornot2,[inp],[outp2,n_out2]).then(tot_prob,None,None)
    #btn.click(aiornot3,[inp],[outp3,n_out3])
    #btn.click(aiornot4,[inp],[outp4,n_out4])
    #btn.click(aiornot5,[inp],[outp5,n_out5])    
app.queue(concurrency_count=20).launch()