update colorize mask

app.py

@@ -13,9 +13,10 @@ def visualize_instance_seg_mask(mask):
     image = np.zeros((mask.shape[0], mask.shape[1], 3))
     labels = np.unique(mask)
     label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
-    for i in range(image.shape[0]):
-      for j in range(image.shape[1]):
-        image[i, j, :] = label2color[mask[i, j]]
+    
+    for id in labels:
+        image[mask == id] = label2color[id]
+
     image = image / 255
     return image