Update app.py

app.py

@@ -1,6 +1,8 @@
 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
 from PIL import Image
@@ -8,20 +10,64 @@ import urllib.request
 import uuid
 uid=uuid.uuid4()
 
+models=[
+    "Nahrawy/AIorNot",
+    "umm-maybe/AI-image-detector",
+    "arnolfokam/ai-generated-image-detector",
+]
 
+pipe0 = pipeline("image-classification", f"{models[0]}")
+pipe1 = pipeline("image-classification", f"{models[1]}")
+pipe2 = pipeline("image-classification", f"{models[2]}")
+ 
+
+fin_sum=[] 
+def image_classifier0(image):
+    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) 
+    fin_sum.append(results)
+    return results
+def image_classifier1(image):
+    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)
+    fin_sum.append(results)
+    return results
+def image_classifier2(image):
+    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) 
+    fin_sum.append(results)
+    return results
 
 def softmax(vector):
  e = exp(vector)
  return e / e.sum()
 
-    
-models=[
-    "Nahrawy/AIorNot",
-    "umm-maybe/AI-image-detector",
-    "arnolfokam/ai-generated-image-detector",
+     
 
-]
-fin_sum=[]   
 def aiornot0(image):    
     labels = ["Real", "AI"]
     mod=models[0]
@@ -39,8 +85,8 @@ def aiornot0(image):
     <h1>This image is likely: {label}</h1><br><h3>
  
     Probabilites:<br>
-    Real: {px[0][0]}<br>
-    AI: {px[1][0]}"""
+    Real: {px[1][0]}<br>
+    AI: {px[0][0]}"""
     results = {}
     for idx,result in enumerate(px):
         results[labels[idx]] = px[idx][0]
@@ -48,7 +94,7 @@ def aiornot0(image):
     fin_sum.append(results)
     return gr.HTML.update(html_out),results
 def aiornot1(image):    
-    labels = ["Real", "AI"]
+    labels = ["AI", "Real"]
     mod=models[1]
     feature_extractor1 = AutoFeatureExtractor.from_pretrained(mod)
     model1 = AutoModelForImageClassification.from_pretrained(mod)
@@ -64,8 +110,8 @@ def aiornot1(image):
     <h1>This image is likely: {label}</h1><br><h3>
   
     Probabilites:<br>
-    Real: {px[0][0]}<br>
-    AI: {px[1][0]}"""
+    Real: {px[1][0]}<br>
+    AI: {px[0][0]}"""
     results = {}
     for idx,result in enumerate(px):
         results[labels[idx]] = px[idx][0]
@@ -73,9 +119,9 @@ def aiornot1(image):
     fin_sum.append(results)
     return gr.HTML.update(html_out),results    
 def aiornot2(image):    
-    labels = ["AI", "Real"]
+    labels = ["Real", "AI"]
     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():
@@ -89,8 +135,8 @@ def aiornot2(image):
     <h1>This image is likely: {label}</h1><br><h3>
   
     Probabilites:<br>
-    Real: {px[1][0]}<br>
-    AI: {px[0][0]}"""
+    Real: {px[0][0]}<br>
+    AI: {px[1][0]}"""
 
     results = {}
     for idx,result in enumerate(px):
@@ -114,15 +160,15 @@ def load_url(url):
 
 def tot_prob():
     try:
-        fin_out = fin_sum[0]["Real"]+fin_sum[1]["Real"]+fin_sum[2]["Real"]
-        fin_out = fin_out/3
+        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_sum.clear()
-        print (fin_out)
+        #print (fin_out)
         return out
     except Exception as e:
         pass
@@ -134,7 +180,7 @@ def fin_clear():
 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='numpy')
+        inp = gr.Image(type='pil')
         in_url=gr.Textbox(label="Image URL")
         with gr.Row():
             load_btn=gr.Button("Load URL")
@@ -164,11 +210,36 @@ with gr.Blocks() as app:
             with gr.Box():
                 n_out2=gr.Label(label="Output")
                 outp2 = gr.HTML("""""")    
-          
-    btn.click(fin_clear,None,fin)
+        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)            
+    
+    def upd(image):
+        rand_im = uuid.uuid4()
+        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_url 
+    #inp.change(upd,inp,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)
-    btn.click(aiornot1,[inp],[outp1,n_out1]).then(tot_prob,None,fin)
-    btn.click(aiornot2,[inp],[outp2,n_out2]).then(tot_prob,None,fin)
     
-app.queue(concurrency_count=20).launch()
+    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)
+
+app.queue(concurrency_count=60).launch()