Update app.py

app.py

@@ -87,32 +87,30 @@ def sepia(input_img):
     )  # We reverse the shape of `image` because `image.size` returns width and height.
     seg = tf.math.argmax(logits, axis=-1)[0]
 
-    color_seg = np.zeros(
-        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
-    )  # height, width, 3
-    for label, color in enumerate(colormap):
-        color_seg[seg.numpy() == label, :] = color
-
-    # Show image + mask
-    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
-    pred_img = pred_img.astype(np.uint8)
-
-    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.figure(figsize=(12, 6))
+    plt.barh(y_pos, probabilities.numpy().mean(axis=(0, 1)), align='center')
+    plt.yticks(y_pos, class_names)
     plt.xlabel('Probability')
     plt.title('Class Probabilities')
 
-    return plt.gcf()
+    # Obtain colored segmentation mask
+    color_seg = label_to_color_image(seg.numpy())
+
+    # Show image + mask + probability bar plot
+    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
+    pred_img = pred_img.astype(np.uint8)
+
+    fig, ax = plt.subplots(1, 2, figsize=(18, 9))
+    ax[0].imshow(pred_img)
+    ax[0].axis('off')
+    ax[1].imshow(plt.gca().get_children()[1].get_children()[0].get_children()[2].get_children()[0].get_children()[0].get_children()[0].get_children()[1].get_children()[1].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[1].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[1].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].get_children()[0].
+
 
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(400, 600)),