Update app.py

app.py

@@ -87,30 +87,18 @@ 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]
 
-    # Obtain probabilities
-    probabilities = tf.nn.softmax(logits, axis=-1)
-
-    # Visualize probabilities as bar plot
-    class_names = labels_list
-    y_pos = np.arange(len(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')
-
-    # Obtain colored segmentation mask
-    color_seg = label_to_color_image(seg.numpy())
-
-    # Show image + mask + probability bar plot
+    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, 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].
-
+    fig = draw_plot(pred_img, seg)
+    return fig
 
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(400, 600)),