Update app.py

app.py

@@ -1,257 +1,216 @@
 import gradio as gr
 import torch
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
-#from transformers import pipeline
 import os 
 from numpy import exp
-import pandas as  pd
+import pandas as pd
 from PIL import Image
 import urllib.request 
 import uuid
-uid=uuid.uuid4()
+uid = uuid.uuid4()
 
-models=[
-    "cmckinle/sdxl-flux-detector",
+# Reordered models as requested
+models = [
     "umm-maybe/AI-image-detector",
     "Organika/sdxl-detector",
-    #"arnolfokam/ai-generated-image-detector",
+    "cmckinle/sdxl-flux-detector",
 ]
 
 pipe0 = pipeline("image-classification", f"{models[0]}")
 pipe1 = pipeline("image-classification", f"{models[1]}")
 pipe2 = pipeline("image-classification", f"{models[2]}")
- 
 
-fin_sum=[] 
+fin_sum = []
+
+def softmax(vector):
+    e = exp(vector - vector.max())  # for numerical stability
+    return e / e.sum()
+
 def image_classifier0(image):
-    labels = ["AI","Real"]
+    labels = ["AI", "Real"]
     outputs = pipe0(image)
     results = {}
-    result_test={}
-    for idx,result in enumerate(outputs):
-        results[labels[idx]] = outputs[idx]['score']
-    #print (result_test)
-    #for result in outputs:
-    #    results[result['label']] = result['score']    
-    #print (results) 
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
     fin_sum.append(results)
     return results
+
 def image_classifier1(image):
-    labels = ["AI","Real"]
+    labels = ["AI", "Real"]
     outputs = pipe1(image)
     results = {}
-    result_test={}
-    for idx,result in enumerate(outputs):
-        results[labels[idx]] = outputs[idx]['score']
-    #print (result_test)    
-    #for result in outputs:
-    #    results[result['label']] = result['score']
-    #print (results)
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
     fin_sum.append(results)
     return results
+
 def image_classifier2(image):
-    labels = ["AI","Real"]
+    labels = ["AI", "Real"]
     outputs = pipe2(image)
     results = {}
-    result_test={}
-    for idx,result in enumerate(outputs):
-        results[labels[idx]] = outputs[idx]['score']
-    #print (result_test)    
-    #for result in outputs:
-    #    results[result['label']] = result['score']
-    #print (results) 
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
     fin_sum.append(results)
     return results
 
-def softmax(vector):
- e = exp(vector)
- return e / e.sum()
-
-     
-
-def aiornot0(image):    
+def aiornot0(image):
     labels = ["AI", "Real"]
-    mod=models[0]
+    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)
+        probability = softmax(logits)  # Apply softmax on 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']
+    Probabilities:<br>
+    Real: {float(px[1][0])}<br>
+    AI: {float(px[0][0])}"""
+    
+    results = {
+        "Real": float(px[1][0]),
+        "AI": float(px[0][0])
+    }
     fin_sum.append(results)
-    return gr.HTML.update(html_out),results
-def aiornot1(image):    
+    return gr.HTML.update(html_out), results
+
+def aiornot1(image):
     labels = ["AI", "Real"]
-    mod=models[1]
+    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)
+        probability = softmax(logits)  # Apply softmax on 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']
+    Probabilities:<br>
+    Real: {float(px[1][0])}<br>
+    AI: {float(px[0][0])}"""
+    
+    results = {
+        "Real": float(px[1][0]),
+        "AI": float(px[0][0])
+    }
     fin_sum.append(results)
-    return gr.HTML.update(html_out),results    
-def aiornot2(image):    
-    labels = ["Real", "AI"]
-    mod=models[2]
+    return gr.HTML.update(html_out), results
+
+def aiornot2(image):
+    labels = ["AI", "Real"]
+    mod = models[2]
     feature_extractor2 = AutoFeatureExtractor.from_pretrained(mod)
-    #feature_extractor2 = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
     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)
+        probability = softmax(logits)  # Apply softmax on 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']
+    Probabilities:<br>
+    Real: {float(px[1][0])}<br>
+    AI: {float(px[0][0])}"""
+    
+    results = {
+        "Real": float(px[1][0]),
+        "AI": float(px[0][0])
+    }
     fin_sum.append(results)
-        
-    return gr.HTML.update(html_out),results
+    return gr.HTML.update(html_out), results
 
 def load_url(url):
     try:
-        urllib.request.urlretrieve( 
-            f'{url}', 
-            f"{uid}tmp_im.png")         
+        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
+        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"]+fin_sum[3]["Real"]+fin_sum[4]["Real"]+fin_sum[5]["Real"]
-        fin_out = fin_out/6
-        fin_sub = 1-fin_out
-        out={
-            "Real":f"{fin_out}",
-            "AI":f"{fin_sub}"
+        fin_out = sum([result["Real"] for result in fin_sum]) / len(fin_sum)
+        fin_sub = 1 - fin_out
+        out = {
+            "Real": f"{fin_out}",
+            "AI": f"{fin_sub}"
         }
-        #fin_sum.clear()
-        #print (fin_out)
         return out
     except Exception as e:
-        pass
-        print (e)
+        print(e)
         return None
+
 def fin_clear():
     fin_sum.clear()
     return None
 
 def upd(image):
-    print (image)
     rand_im = uuid.uuid4()
     image.save(f"{rand_im}-vid_tmp_proc.png")
     out = Image.open(f"{rand_im}-vid_tmp_proc.png")
+    return out
 
-    #image.save(f"{rand_im}-vid_tmp_proc.png")
-    #out = os.path.abspath(f"{rand_im}-vid_tmp_proc.png")
-    #out_url = f'https://omnibus_AI_or_Not_dev.hf.space/file={out}'
-    #out_url = f"{rand_im}-vid_tmp_proc.png"
-    return out 
-
-    
 with gr.Blocks() as app:
     gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)""")
     with gr.Column():
         inp = gr.Image(type='pil')
-        in_url=gr.Textbox(label="Image URL")
+        in_url = gr.Textbox(label="Image URL")
         with gr.Row():
-            load_btn=gr.Button("Load URL")
+            load_btn = gr.Button("Load URL")
             btn = gr.Button("Detect AI")
         mes = gr.HTML("""""")
-    with gr.Group():    
+
+    with gr.Group():
         with gr.Row():
-            fin=gr.Label(label="Final Probability")       
+            fin = gr.Label(label="Final Probability", visible=False)
         with gr.Row():
+            # Updated model names
             with gr.Box():
-                lab0 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[0]}'>{models[0]}</a></b>""")
+                lab0 = gr.HTML(f"""<b>Testing on Original Model: <a href='https://huggingface.co/{models[0]}'>{models[0]}</a></b>""")
                 nun0 = gr.HTML("""""")
             with gr.Box():
-                lab1 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[1]}'>{models[1]}</a></b>""")
+                lab1 = gr.HTML(f"""<b>Testing on SDXL Fine Tuned Model: <a href='https://huggingface.co/{models[1]}'>{models[1]}</a></b>""")
                 nun1 = gr.HTML("""""")
             with gr.Box():
-                lab2 = gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{models[2]}'>{models[2]}</a></b>""")
+                lab2 = gr.HTML(f"""<b>Testing on SDXL and Flux Fine Tuned Model: <a href='https://huggingface.co/{models[2]}'>{models[2]}</a></b>""")
                 nun2 = gr.HTML("""""")
-                
         with gr.Row():
             with gr.Box():
-                n_out0=gr.Label(label="Output")
+                n_out0 = gr.Label(label="Output")
                 outp0 = gr.HTML("""""")
             with gr.Box():
-                n_out1=gr.Label(label="Output")
+                n_out1 = gr.Label(label="Output")
                 outp1 = gr.HTML("""""")
             with gr.Box():
-                n_out2=gr.Label(label="Output")
-                outp2 = gr.HTML("""""")    
-        with gr.Row():
-            with gr.Box():
-                n_out3=gr.Label(label="Output")
-                outp3 = gr.HTML("""""")
-            with gr.Box():
-                n_out4=gr.Label(label="Output")
-                outp4 = gr.HTML("""""")
-            with gr.Box():
-                n_out5=gr.Label(label="Output")
-                outp5 = gr.HTML("""""")    
-    hid_box=gr.Textbox(visible=False)  
-    hid_im = gr.Image(type="pil",visible=False)
-    def echo(inp):
-        return inp
+                n_out2 = gr.Label(label="Output")
+                outp2 = gr.HTML("""""")
 
-    #inp.change(echo,inp,hid_im).then(upd,hid_im,inp)
-    
-    btn.click(fin_clear,None,fin,show_progress=False)
-    load_btn.click(load_url,in_url,[inp,mes])
-    
-    btn.click(aiornot0,[inp],[outp0,n_out0]).then(tot_prob,None,fin,show_progress=False)
-    btn.click(aiornot1,[inp],[outp1,n_out1]).then(tot_prob,None,fin,show_progress=False)
-    btn.click(aiornot2,[inp],[outp2,n_out2]).then(tot_prob,None,fin,show_progress=False)
-    
-    btn.click(image_classifier0,[inp],[n_out3]).then(tot_prob,None,fin,show_progress=False)
-    btn.click(image_classifier1,[inp],[n_out4]).then(tot_prob,None,fin,show_progress=False)
-    btn.click(image_classifier2,[inp],[n_out5]).then(tot_prob,None,fin,show_progress=False)
+    btn.click(fin_clear, None, fin, show_progress=False)
+    load_btn.click(load_url, in_url, [inp, mes])
+
+    btn.click(aiornot0, [inp], [outp0, n_out0]).then(tot_prob, None, fin, show_progress=False)
+    btn.click(aiornot1, [inp], [outp1, n_out1]).then(tot_prob, None, fin, show_progress=False)
+    btn.click(aiornot2, [inp], [outp2, n_out2]).then(tot_prob, None, fin, show_progress=False)
+
+    btn.click(image_classifier0, [inp], [n_out0]).then(tot_prob, None, fin, show_progress=False)
+    btn.click(image_classifier1, [inp], [n_out1]).then(tot_prob, None, fin, show_progress=False)
+    btn.click(image_classifier2, [inp], [n_out2]).then(tot_prob, None, fin, show_progress=False)
 
-app.launch(show_api=False,max_threads=24)
+app.launch(show_api=False, max_threads=24)