Update app.py

app.py

@@ -100,6 +100,20 @@ def sepia(input_img):
     fig = draw_plot(pred_img, seg)
     return fig
 
+    # Obtain probabilities
+    probabilities = tf.nn.softmax(logits, axis=-1)
+
+    # Visualize probabilities as bar plot
+    plt.figure(figsize=(12,6))
+    class_names = labels_list
+    y_pos = np.arange(len(class_names))
+    plt.bart(y_pos, probabilities.numpy().mean(axis=(0, 1)), align='center')
+    plt.ytickes(y_pos, class_names)
+    plt.xlabel('Probability')
+    plt.title('Class Probabilities')
+
+    return plt.gcf()
+
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(400, 600)),
                     outputs=['plot'],