Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import torch
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
 from numpy import exp
- 
+import pandas as  pd
+
 def softmax(vector):
  e = exp(vector)
  return e / e.sum()
@@ -10,68 +11,107 @@ def softmax(vector):
     
 models=[
     "Nahrawy/AIorNot",
-    "arnolfokam/ai-generated-image-detector",
     "umm-maybe/AI-image-detector",
-]
+    "arnolfokam/ai-generated-image-detector",
 
+]
+    
 def aiornot0(image):    
     labels = ["Real", "AI"]
     mod=models[0]
-    feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
-    model = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor(image, return_tensors="pt")
+    feature_extractor0 = AutoFeatureExtractor.from_pretrained(mod)
+    model0 = AutoModelForImageClassification.from_pretrained(mod)
+    input = feature_extractor0(image, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(**input)
-        print (outputs)
+        outputs = model0(**input)
         logits = outputs.logits
-        print (logits)
         probability = softmax(logits)
-        print(probability)
-
+        px = pd.DataFrame(probability.numpy())
     prediction = logits.argmax(-1).item()
-    label = labels[prediction] 
-    return label
+    label = labels[prediction]
+    html_out = f"""
+    <h1>This image is likely: {label}</h1><br><h3>
+    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
+    <br>    
+    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 aiornot1(image):    
     labels = ["Real", "AI"]
     mod=models[1]
-    feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
-    model = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor(image, return_tensors="pt")
+    feature_extractor1 = AutoFeatureExtractor.from_pretrained(mod)
+    model1 = AutoModelForImageClassification.from_pretrained(mod)
+    input = feature_extractor1(image, return_tensors="pt")
     with torch.no_grad():
-      outputs = model(**input)
-      print (outputs)
-      logits = outputs.logits
-      print (logits)
+        outputs = model1(**input)
+        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"""
+    <h1>This image is likely: {label}</h1><br><h3>
+    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
+    <br>    
+    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 aiornot2(image):    
-    labels = ["Real", "AI"]
+    labels = ["AI", "Real"]
     mod=models[2]
-    feature_extractor = AutoFeatureExtractor.from_pretrained(mod)
-    model = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor(image, return_tensors="pt")
+    feature_extractor2 = AutoFeatureExtractor.from_pretrained(mod)
+    model2 = AutoModelForImageClassification.from_pretrained(mod)
+    input = feature_extractor2(image, return_tensors="pt")
     with torch.no_grad():
-      outputs = model(**input)
-      print (outputs)
-      logits = outputs.logits
-      print (logits)
+        outputs = model2(**input)
+        logits = outputs.logits
+        probability = softmax(logits)
+        px = pd.DataFrame(probability.numpy())
     prediction = logits.argmax(-1).item()
-    label = labels[prediction] 
-    return label    
-with gr.Blocks() as app:
-    with gr.Row():
-        with gr.Column():
-            inp = gr.Image()
-            mod_choose=gr.Number(value=0)
-            btn = gr.Button()
-            
-        with gr.Column():
-            outp0 = gr.Textbox(label=f'{models[0]}')
-            outp1 = gr.Textbox(label=f'{models[1]}')
-            outp2 = gr.Textbox(label=f'{models[2]}')
-    btn.click(aiornot0,[inp],outp0)
-    btn.click(aiornot1,[inp],outp1)
-    btn.click(aiornot2,[inp],outp2)
+    label = labels[prediction]
+    html_out = f"""
+    <h1>This image is likely: {label}</h1><br><h3>
+    Model used: <a href='https://huggingface.co/{mod}'>{mod}</a><br>    
+    <br>    
+    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
     
-app.launch()
+with gr.Blocks() as app:
+    with gr.Column():
+        inp = gr.Pil()
+        btn = gr.Button()
+    with gr.Group():        
+        with gr.Row():
+            with gr.Box():
+                lab0 = gr.HTML(f"""<b>Testing on Model: {models[0]}</b>""")
+                outp0 = gr.HTML("""""")
+                n_out0=gr.Label(label="Output")
+            with gr.Box():
+                lab1 = gr.HTML(f"""<b>Testing on Model: {models[1]}</b>""")
+                outp1 = gr.HTML("""""")
+                n_out1=gr.Label(label="Output")
+            with gr.Box():
+                lab2 = gr.HTML(f"""<b>Testing on Model: {models[2]}</b>""")
+                outp2 = gr.HTML("""""")            
+                n_out2=gr.Label(label="Output")
+    btn.click(aiornot0,[inp],[outp0,n_out0])
+    btn.click(aiornot1,[inp],[outp1,n_out1])
+    btn.click(aiornot2,[inp],[outp2,n_out2])
+app.launch(enable_queue=False)