shelbao commited on
Commit
2a1c65a
·
verified ·
1 Parent(s): dacc836

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
2
+ from PIL import Image
3
+ import requests
4
+ import torch
5
+
6
+ # Load Mask2Former trained on COCO instance segmentation dataset
7
+ image_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-coco-instance")
8
+ model = Mask2FormerForUniversalSegmentation.from_pretrained(
9
+ "facebook/mask2former-swin-small-coco-instance"
10
+ )
11
+
12
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
13
+ image = Image.open(requests.get(url, stream=True).raw)
14
+ inputs = image_processor(image, return_tensors="pt")
15
+
16
+ with torch.no_grad():
17
+ outputs = model(**inputs)
18
+
19
+ # Model predicts class_queries_logits of shape `(batch_size, num_queries)`
20
+ # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
21
+ class_queries_logits = outputs.class_queries_logits
22
+ masks_queries_logits = outputs.masks_queries_logits
23
+
24
+ # Perform post-processing to get instance segmentation map
25
+ pred_instance_map = image_processor.post_process_semantic_segmentation(
26
+ outputs, target_sizes=[image.size[::-1]]
27
+ )[0]
28
+ print(pred_instance_map.shape)