a

app.py

@@ -53,11 +53,7 @@ def greet(input_img):
     )
     seg = tf.math.argmax(logits_tf, 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
+    color_seg = label_to_color_image(seg.numpy())
 
     # Resize color_seg to match the shape of input_img
     color_seg_resized = tf.image.resize(color_seg, (input_img.shape[0], input_img.shape[1]))
@@ -67,10 +63,9 @@ def greet(input_img):
     # Convert pred_img to NumPy array and then change data type
     pred_img = np.array(pred_img).astype(np.uint8)
 
-    fig = draw_plot(pred_img, seg)
+    fig = draw_plot(pred_img, seg.numpy())
     return fig
 
-
 def draw_plot(pred_img, seg):
     fig = plt.figure(figsize=(20, 15))
     grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
@@ -78,6 +73,7 @@ def draw_plot(pred_img, seg):
     plt.subplot(grid_spec[0])
     plt.imshow(pred_img)
     plt.axis("off")
+
     LABEL_NAMES = np.asarray(labels_list)
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
     FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)