Update app.py

app.py

@@ -19,16 +19,9 @@ 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=
+fin_sum=[]   
 def aiornot0(image):    
     labels = ["Real", "AI"]
     mod=models[0]
@@ -52,7 +45,6 @@ def aiornot0(image):
     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):    
@@ -78,8 +70,7 @@ def aiornot1(image):
     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)        
+    fin_sum.append(results)
     return gr.HTML.update(html_out),results    
 def aiornot2(image):    
     labels = ["AI", "Real"]
@@ -101,85 +92,14 @@ def aiornot2(image):
     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']
+    fin_sum.append(results)
+        
     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( 
@@ -201,45 +121,53 @@ def tot_prob():
             "Real":f"{fin_out}",
             "AI":f"{fin_sub}"
         }
+        #fin_sum.clear()
         print (fin_out)
         return out
     except Exception as e:
-        
         pass
-        print (f'ERROR :: {e}')
-
-def clear_fin():
+        print (e)
+def fin_clear():
     fin_sum.clear()
     
 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("""""")
+    gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)""")
+    with gr.Column():
         inp = gr.Pil()
-    with gr.Group():        
+        in_url=gr.Textbox(label="Image URL")
+        with gr.Row():
+            load_btn=gr.Button("Load URL")
+            btn = gr.Button("Detect AI")
+        mes = gr.HTML("""""")
+    with gr.Group():    
+        with gr.Row():
+            fin=gr.Label(label="Final Probability")       
         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>""")
+                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>""")
+                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>""")
+                nun2 = gr.HTML("""""")
+                
+        with gr.Row():
+            with gr.Box():
                 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():
-            fin=gr.Label(label="Final Probability")        
-    fin.change(clear_fin,None,None)
+          
+    fin.change(fin_clear,None,None)
     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()