Update app.py

app.py

@@ -23,25 +23,19 @@ def aiornot0(image):
     input = feature_extractor0(image, return_tensors="pt")
     with torch.no_grad():
         outputs = model0(**input)
-        print (outputs)
         logits = outputs.logits
-        print (logits)
         probability = softmax(logits)
-        print(f'PROBABILITY ::: {probability}')
-        print(probability[0][0])
         px = pd.DataFrame(probability.numpy())
-        print(px)
-        
     prediction = logits.argmax(-1).item()
     label = labels[prediction]
-    
     html_out = f"""
-    <h3>Model used: <a href src='https://huggingface.co/{mod}'>{mod}</a><br>
-    This image is likely: {label}<br>
-    Probabilites<br>
-    AI: {px[0][0]}<br>
-    Real: {px[1][0]}"""
-    
+    <h3>
+    <b>This image is likely: {label}</b><br>
+    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
+    <br>    
+    Probabilites:<br>
+    Real: {px[0][0]}<br>
+    AI: {px[1][0]}"""
     return gr.HTML.update(html_out)
 def aiornot1(image):    
     labels = ["Real", "AI"]
@@ -51,12 +45,20 @@ def aiornot1(image):
     input = feature_extractor1(image, return_tensors="pt")
     with torch.no_grad():
       outputs = model1(**input)
-      print (outputs)
-      logits = outputs.logits
-      print (logits)
+        logits = outputs.logits
+        probability = softmax(logits)
+        px = pd.DataFrame(probability.numpy())
     prediction = logits.argmax(-1).item()
-    label = labels[prediction] 
-    return label
+    label = labels[prediction]
+    html_out = f"""
+    <h3>
+    <b>This image is likely: {label}</b><br>
+    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
+    <br>    
+    Probabilites:<br>
+    Real: {px[0][0]}<br>
+    AI: {px[1][0]}"""
+    return gr.HTML.update(html_out)
 def aiornot2(image):    
     labels = ["Real", "AI"]
     mod=models[2]
@@ -65,12 +67,20 @@ def aiornot2(image):
     input = feature_extractor2(image, return_tensors="pt")
     with torch.no_grad():
       outputs = model2(**input)
-      print (outputs)
-      logits = outputs.logits
-      print (logits)
+        logits = outputs.logits
+        probability = softmax(logits)
+        px = pd.DataFrame(probability.numpy())
     prediction = logits.argmax(-1).item()
-    label = labels[prediction] 
-    return label    
+    label = labels[prediction]
+    html_out = f"""
+    <h3>
+    <b>This image is likely: {label}</b><br>
+    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
+    <br>    
+    Probabilites:<br>
+    Real: {px[0][0]}<br>
+    AI: {px[1][0]}"""
+    return gr.HTML.update(html_out)
 with gr.Blocks() as app:
     with gr.Row():
         with gr.Column():
@@ -80,10 +90,12 @@ with gr.Blocks() as app:
             
         with gr.Column():
             #outp0 = gr.Textbox(label=f'{models[0]}')
-            lab0 = gr.HTML(f"""Testing on Model {models[0]}""")
+            lab0 = gr.HTML(f"""Testing on Model: {models[0]}""")
             outp0 = gr.HTML("""""")
-            outp1 = gr.Textbox(label=f'{models[1]}')
-            outp2 = gr.Textbox(label=f'{models[2]}')
+            lab1 = gr.HTML(f"""Testing on Model: {models[1]}""")
+            outp1 = gr.HTML("""""")
+            lab2 = gr.HTML(f"""Testing on Model: {models[2]}""")
+            outp2 = gr.HTML("""""")            
     btn.click(aiornot0,[inp],outp0)
     btn.click(aiornot1,[inp],outp1)
     btn.click(aiornot2,[inp],outp2)