Upload with huggingface_hub

DESCRIPTION.md

@@ -0,0 +1 @@
+Image segmentation using DETR. Takes in both an inputu image and the desired confidence, and returns a segmented image.

README.md

@@ -1,12 +1,11 @@
+
 ---
-title: Image Segmentation Main
-emoji: 💻
-colorFrom: pink
-colorTo: pink
+title: image_segmentation_main 
+emoji: 🔥
+colorFrom: indigo
+colorTo: indigo
 sdk: gradio
 sdk_version: 3.6
-app_file: app.py
+app_file: run.py
 pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+torch
+scipy
+numpyhttps://gradio-main-build.s3.amazonaws.com/c3bec6153737855510542e8154391f328ac72606/gradio-3.6-py3-none-any.whl

run.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import torch
+import random
+import numpy as np
+from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
+
+device = torch.device("cpu")
+model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade").to(device)
+model.eval()
+preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")
+
+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]]
+    image = image / 255
+    return image
+
+def query_image(img):
+    target_size = (img.shape[0], img.shape[1])
+    inputs = preprocessor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    outputs.class_queries_logits = outputs.class_queries_logits.cpu()
+    outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()
+    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()
+    results = torch.argmax(results, dim=0).numpy()
+    results = visualize_instance_seg_mask(results)
+    return results
+
+demo = gr.Interface(
+    query_image, 
+    inputs=[gr.Image()], 
+    outputs="image",
+    title="MaskFormer Demo",
+    examples=[["example_2.png"]]
+)
+
+demo.launch()