using new labels

app.py

@@ -9,7 +9,7 @@ ade_mean=[0.485, 0.456, 0.406]
 ade_std=[0.229, 0.224, 0.225]
 
 palette = [
-    [120, 120, 120], [4, 200, 4], [180, 120, 120], [6, 230, 230],
+    [120, 120, 120], [4, 200, 4], [4, 4, 250], [6, 230, 230],
     [80, 50, 50], [120, 120, 80], [140, 140, 140], [204, 5, 255]
 ]
 
@@ -36,6 +36,7 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
     img_out = np.zeros((mask.shape[0], mask.shape[1], 3))
     image_total_pixels = mask.shape[0] * mask.shape[1]
     label_ids = np.unique(mask)
+    vegetation_labels = ["vegetation"]
 
     id2color = {id: palette[id] for id in label_ids}
     id2count = {id: 0 for id in label_ids}
@@ -51,7 +52,7 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
         f"{id2label[id]}",
         f"{(100 * id2count[id] / image_total_pixels):.2f} %",
         f"{np.sqrt(id2count[id] / image_total_pixels):.2f} m"
-        ] for id in label_ids if 'tree' in id2label[id]]
+        ] for id in label_ids if id2label[id] in vegetation_labels]
 
     return image_res, dataframe