adds examples

README.md

@@ -4,7 +4,7 @@ emoji: 🥦
 colorFrom: pink
 colorTo: indigo
 sdk: gradio
-sdk_version: 3.16.1
+sdk_version: 3.16.2
 app_file: app.py
 models:
     - "facebook/maskformer-swin-large-coco"

app.py

@@ -1,5 +1,7 @@
+import glob
 import gradio as gr
 import numpy as np
+from PIL import Image
 from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
 
 
@@ -8,6 +10,8 @@ from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmen
 feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
 model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
 
+example_images = sorted(glob.glob('examples/map*.jpg'))
+
 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]
@@ -28,7 +32,7 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
         img_out[i, j, :] = id2color[mask[i, j]]
         id2count[mask[i, j]] = id2count[mask[i, j]] + 1
 
-    image_res = (0.5 * img_in + 0.5 * img_out) / 255
+    image_res = (0.5 * img_in + 0.5 * img_out).astype(np.uint8)
 
     vegetation_count = sum([id2count[id] for id in label_ids if id2label[id] in vegetation_labels])
 
@@ -47,28 +51,39 @@ def visualize_instance_seg_mask(img_in, mask, id2label):
         f"{(100 * vegetation_count / image_total_pixels):.2f} %",
         f"{np.sqrt(vegetation_count / image_total_pixels):.2f} m"]]
 
-    return image_res, dataframe_vegetation_total
+    dataframe = dataframe_vegetation_total
+    if len(dataframe) < 1:
+        dataframe = [[
+            f"",
+            f"{(0):.2f} %",
+            f"{(0):.2f} m"
+        ]]
+
+    return image_res, dataframe
 
 
-def query_image(img):
+def query_image(image_path):
+    img = np.array(Image.open(image_path))
     img_size = (img.shape[0], img.shape[1])
     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(img, results.numpy(), model.config.id2label)
-    return results
+    mask_img, dataframe = visualize_instance_seg_mask(img, results.numpy(), model.config.id2label)
+    return mask_img, dataframe
 
 
 demo = gr.Interface(
     query_image,
-    inputs=[gr.Image(label="Input Image")],
+    inputs=[gr.Image(type="filepath", label="Input Image")],
     outputs=[
         gr.Image(label="Vegetation"),
         gr.DataFrame(label="Info", headers=["Object Label", "Pixel Percent", "Square Length"])
     ],
     title="Maskformer (large-coco)",
     allow_flagging="never",
-    analytics_enabled=None
+    analytics_enabled=None,
+    examples=example_images,
+    cache_examples=True
 )
 
 demo.launch(show_api=False)

requirements.txt

@@ -1,3 +1,4 @@
-transformers
+Pillow
+scipy
 torch
-scipy
+transformers