Update app.py

app.py

@@ -26,33 +26,40 @@ 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)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+    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))}
     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)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+    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))}
     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)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+    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))}
     fin_sum.append(results)
     return results
 
+# AI or Not function for Model 0
 def aiornot0(image):
     labels = ["AI", "Real"]
     mod = models[0]
@@ -80,6 +87,7 @@ 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]
@@ -107,6 +115,7 @@ 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]
@@ -134,6 +143,7 @@ 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")
@@ -144,6 +154,7 @@ 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)
@@ -157,10 +168,12 @@ 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")
@@ -181,7 +194,6 @@ 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("""""")
@@ -205,12 +217,9 @@ 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)