Adding probs

app.py

@@ -5,6 +5,7 @@ Repo.clone_from("https://github.com/dimitreOliveira/hub.git", "./hub")
 sys.path.append("/hub")
 
 import gradio as gr
+import tensorflow as tf
 from hub.tensorflow_hub.hf_utils import pull_from_hub
 
 import requests
@@ -20,26 +21,12 @@ def preprocess(image):
 
 def postprocess(prediction):
     return {labels[i]: float(prediction[i]) for i in range(len(labels))}
-    # return {labels[i]: 0 for i in range(len(labels))}
 
 def predict_fn(image):
     image = preprocess(image)
-    prediction = model(image)[0].numpy()
-    print('****************')
-    print(prediction)
-    try:
-        print("default")
-        print({labels[i]: float(prediction[i]) for i in range(len(labels))})
-    except:
-        print("default gives error")
-    print('****************')
-    print(list(prediction))
-    try:
-        print("list")
-        print({labels[i]: list(prediction)[i] for i in range(len(labels))})
-    except:
-        print("list gives error")
-    scores = postprocess(prediction)
+    logits = model(image)
+    probs = tf.nn.softmax(logits, axis=1)[0].numpy()
+    scores = postprocess(probs)
     return scores
 
 iface = gr.Interface(fn=predict_fn,