a

app.py

@@ -79,11 +79,22 @@ def draw_plot(pred_img, seg):
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
     FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
 
+    # Limit unique_labels to be within the range of colormap
     unique_labels = np.unique(seg.astype("uint8"))
+    unique_labels = unique_labels[unique_labels < len(colormap)]
+
     ax = plt.subplot(grid_spec[1])
-    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
-    ax.yaxis.tick_right()
-    plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
+
+    if len(unique_labels) > 0:
+        plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
+        ax.yaxis.tick_right()
+        plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
+    else:
+        # Handle case when there are no unique labels
+        plt.imshow(np.zeros((1, 1, 3), dtype=np.uint8))
+        ax.yaxis.tick_right()
+        plt.yticks([], [])
+
     plt.xticks([], [])
     ax.tick_params(width=0.0, labelsize=25)
     return fig