aliabd HF staff commited on
Commit
d074d90
1 Parent(s): e39cf81

Upload with huggingface_hub

Browse files
Files changed (2) hide show
  1. README.md +1 -2
  2. run.py +42 -0
README.md CHANGED
@@ -6,7 +6,6 @@ colorFrom: indigo
6
  colorTo: indigo
7
  sdk: gradio
8
  sdk_version: 3.4.1
9
-
10
- app_file: app.py
11
  pinned: false
12
  ---
 
6
  colorTo: indigo
7
  sdk: gradio
8
  sdk_version: 3.4.1
9
+ app_file: run.py
 
10
  pinned: false
11
  ---
run.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import random
4
+ import numpy as np
5
+ from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
6
+
7
+ device = torch.device("cpu")
8
+ model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade").to(device)
9
+ model.eval()
10
+ preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")
11
+
12
+ def visualize_instance_seg_mask(mask):
13
+ image = np.zeros((mask.shape[0], mask.shape[1], 3))
14
+ labels = np.unique(mask)
15
+ label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
16
+ for i in range(image.shape[0]):
17
+ for j in range(image.shape[1]):
18
+ image[i, j, :] = label2color[mask[i, j]]
19
+ image = image / 255
20
+ return image
21
+
22
+ def query_image(img):
23
+ target_size = (img.shape[0], img.shape[1])
24
+ inputs = preprocessor(images=img, return_tensors="pt")
25
+ with torch.no_grad():
26
+ outputs = model(**inputs)
27
+ outputs.class_queries_logits = outputs.class_queries_logits.cpu()
28
+ outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()
29
+ results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()
30
+ results = torch.argmax(results, dim=0).numpy()
31
+ results = visualize_instance_seg_mask(results)
32
+ return results
33
+
34
+ demo = gr.Interface(
35
+ query_image,
36
+ inputs=[gr.Image()],
37
+ outputs="image",
38
+ title="MaskFormer Demo",
39
+ examples=[["example_2.png"]]
40
+ )
41
+
42
+ demo.launch()