Views
No views yet

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