blends in and out images

app.py

@@ -8,20 +8,26 @@ model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-s
 # feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
 # model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
 
-def visualize_instance_seg_mask(mask, id2label):
-    image = np.zeros((mask.shape[0], mask.shape[1], 3))
+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)
 
-    id2color = {id: (np.random.randint(0, 2), np.random.randint(0, 256), np.random.randint(0, 256)) for id in label_ids}
+    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]:
+            id_color = (0, 140, 0)
+        return id_color
+
+    id2color = {id: get_color(id) for id in label_ids}
     id2count = {id: 0 for id in label_ids}
 
-    for i in range(image.shape[0]):
-      for j in range(image.shape[1]):
-        image[i, j, :] = id2color[mask[i, j]]
+    for i in range(img_out.shape[0]):
+      for j in range(img_out.shape[1]):
+        img_out[i, j, :] = id2color[mask[i, j]]
         id2count[mask[i, j]] = id2count[mask[i, j]] + 1
 
-    image = image / 255
+    image_res = (0.5 * img_in + 0.5 * img_out) / 255
 
     dataframe = [[
         f"{id2label[id]}",
@@ -29,7 +35,7 @@ def visualize_instance_seg_mask(mask, id2label):
         f"{np.sqrt(id2count[id] / image_total_pixels):.2f} m"
         ] for id in label_ids if 'tree' in id2label[id]]
 
-    return image, dataframe
+    return image_res, dataframe
 
 
 def query_image(img):
@@ -37,14 +43,14 @@ def query_image(img):
     inputs = feature_extractor(images=img, return_tensors="pt")
     outputs = model(**inputs)
     results = feature_extractor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
-    results = visualize_instance_seg_mask(results.numpy(), model.config.id2label)
+    results = visualize_instance_seg_mask(img, results.numpy(), model.config.id2label)
     return results
 
 
 demo = gr.Interface(
     query_image,
-    inputs=[gr.Image()],
-    outputs=["image", gr.DataFrame(headers=None, default_value=None, label="Area Info")],
+    inputs=[gr.Image(label="Input Image")],
+    outputs=[gr.Image(label="Trees"), gr.DataFrame(headers=None, default_value=None, label="Area Info")],
     title="maskformer-swin-large-coco",
     allow_flagging="never",
     analytics_enabled=None