Update app.py

app.py

@@ -1,7 +1,13 @@
 import gradio as gr
 import torch
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
+from numpy import exp
+ 
+def softmax(vector):
+ e = exp(vector)
+ return e / e.sum()
 
+    
 models=[
     "Nahrawy/AIorNot",
     "arnolfokam/ai-generated-image-detector",
@@ -15,10 +21,13 @@ def aiornot0(image):
     model = AutoModelForImageClassification.from_pretrained(mod)
     input = feature_extractor(image, return_tensors="pt")
     with torch.no_grad():
-      outputs = model(**input)
-      print (outputs)
-      logits = outputs.logits
-      print (logits)
+        outputs = model(**input)
+        print (outputs)
+        logits = outputs.logits
+        print (logits)
+        probability = softmax(logits, axis=-1)
+        print(probability)
+
     prediction = logits.argmax(-1).item()
     label = labels[prediction] 
     return label