add tree+grass for vegetation count

app.py

@@ -3,19 +3,20 @@ import numpy as np
 from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
 
 
-feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-coco")
-model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-coco")
-# feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
-# model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
+# feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-coco")
+# model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-coco")
+feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
+model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
 
 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 = ["tree-merged", "grass-merged"]
 
     def get_color(id):
         id_color = (np.random.randint(0, 2), np.random.randint(0, 4), np.random.randint(0, 256))
-        if 'tree' in id2label[id]:
+        if id2label[id] in vegetation_labels:
             id_color = (0, 140, 0)
         return id_color
 
@@ -29,13 +30,24 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
 
     image_res = (0.5 * img_in + 0.5 * img_out) / 255
 
-    dataframe = [[
+    vegetation_count = sum([id2count[id] for id in label_ids if id2label[id] in vegetation_labels])
+
+    dataframe_vegetation_items = [[
+        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 id2label[id] in vegetation_labels]
+    dataframe_all_items = [[
         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]
+    dataframe_vegetation_total = [[
+        f"vegetation",
+        f"{(100 * vegetation_count / image_total_pixels):.2f} %",
+        f"{np.sqrt(vegetation_count / image_total_pixels):.2f} m"]]
 
-    return image_res, dataframe
+    return image_res, dataframe_vegetation_total
 
 
 def query_image(img):
@@ -50,7 +62,7 @@ def query_image(img):
 demo = gr.Interface(
     query_image,
     inputs=[gr.Image(label="Input Image")],
-    outputs=[gr.Image(label="Trees"), gr.DataFrame(headers=None, label="Area Info")],
+    outputs=[gr.Image(label="Vegetation"), gr.DataFrame(headers=None, label="Area Info")],
     title="maskformer-swin-large-coco",
     allow_flagging="never",
     analytics_enabled=None