Update app.py

app.py

@@ -26,40 +26,33 @@ def softmax(vector):
     e = exp(vector - vector.max())  # for numerical stability
     return e / e.sum()
 
-# Image classification function for Model 0
 def image_classifier0(image):
-    fin_sum.clear()  # Clear previous results
     labels = ["AI", "Real"]
     outputs = pipe0(image)
-    scores = [output['score'] for output in outputs]
-    soft_scores = softmax(scores)  # Ensure consistency with softmax
-    results = {labels[i]: float(soft_scores[i]) for i in range(len(labels))}
+    results = {}
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
     fin_sum.append(results)
     return results
 
-# Image classification function for Model 1
 def image_classifier1(image):
-    fin_sum.clear()  # Clear previous results
     labels = ["AI", "Real"]
     outputs = pipe1(image)
-    scores = [output['score'] for output in outputs]
-    soft_scores = softmax(scores)  # Ensure consistency with softmax
-    results = {labels[i]: float(soft_scores[i]) for i in range(len(labels))}
+    results = {}
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
     fin_sum.append(results)
     return results
 
-# Image classification function for Model 2
 def image_classifier2(image):
-    fin_sum.clear()  # Clear previous results
     labels = ["AI", "Real"]
     outputs = pipe2(image)
-    scores = [output['score'] for output in outputs]
-    soft_scores = softmax(scores)  # Ensure consistency with softmax
-    results = {labels[i]: float(soft_scores[i]) for i in range(len(labels))}
+    results = {}
+    for idx, result in enumerate(outputs):
+        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
     fin_sum.append(results)
     return results
 
-# AI or Not function for Model 0
 def aiornot0(image):
     labels = ["AI", "Real"]
     mod = models[0]
@@ -87,7 +80,6 @@ def aiornot0(image):
     fin_sum.append(results)
     return gr.HTML.update(html_out), results
 
-# AI or Not function for Model 1
 def aiornot1(image):
     labels = ["AI", "Real"]
     mod = models[1]
@@ -115,7 +107,6 @@ def aiornot1(image):
     fin_sum.append(results)
     return gr.HTML.update(html_out), results
 
-# AI or Not function for Model 2
 def aiornot2(image):
     labels = ["AI", "Real"]
     mod = models[2]
@@ -143,7 +134,6 @@ def aiornot2(image):
     fin_sum.append(results)
     return gr.HTML.update(html_out), results
 
-# Load URL and return image
 def load_url(url):
     try:
         urllib.request.urlretrieve(f'{url}', f"{uid}tmp_im.png")
@@ -154,7 +144,6 @@ def load_url(url):
         mes = f"Image not Found<br>Error: {e}"
     return image, mes
 
-# Calculate final probabilities
 def tot_prob():
     try:
         fin_out = sum([result["Real"] for result in fin_sum]) / len(fin_sum)
@@ -168,12 +157,10 @@ def tot_prob():
         print(e)
         return None
 
-# Clear the fin_sum list
 def fin_clear():
     fin_sum.clear()
     return None
 
-# Update image
 def upd(image):
     rand_im = uuid.uuid4()
     image.save(f"{rand_im}-vid_tmp_proc.png")
@@ -194,6 +181,7 @@ with gr.Blocks() as app:
         with gr.Row():
             fin = gr.Label(label="Final Probability", visible=False)
         with gr.Row():
+            # Updated model names
             with gr.Box():
                 lab0 = gr.HTML(f"""<b>Testing on Original Model: <a href='https://huggingface.co/{models[0]}'>{models[0]}</a></b>""")
                 nun0 = gr.HTML("""""")
@@ -217,9 +205,12 @@ with gr.Blocks() as app:
     btn.click(fin_clear, None, fin, show_progress=False)
     load_btn.click(load_url, in_url, [inp, mes])
 
-    # Use either the aiornot functions or image_classifier consistently
     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)