Views
No views yet

1import requests
2import torch
3from PIL import Image
4from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
5
6
7# load Mask2Former fine-tuned on COCO instance segmentation
8processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-small-coco-instance")
9model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-small-coco-instance")
10
11url = "http://images.cocodataset.org/val2017/000000039769.jpg"
12image = Image.open(requests.get(url, stream=True).raw)
13inputs = processor(images=image, return_tensors="pt")
14
15with torch.no_grad():
16 outputs = model(**inputs)
17
18# model predicts class_queries_logits of shape `(batch_size, num_queries)`
19# and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
20class_queries_logits = outputs.class_queries_logits
21masks_queries_logits = outputs.masks_queries_logits
22
23# you can pass them to processor for postprocessing
24result = processor.post_process_instance_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
25# we refer to the demo notebooks for visualization (see "Resources" section in the Mask2Former docs)
26predicted_instance_map = result["segmentation"]