Add segmentation model support to from_transformers

[Griffin-Sullivan]

trafficstars

Description

This PR adds support to the from_transformers method to allow getting detections from Segmentation models. At the moment I'm testing the feature by returning the object based on the boxes key being present. Processing segmentation models in transformers doesn't return the boxes key.

List any dependencies that are required for this change.

Type of change

Please delete options that are not relevant.

• [ ] Bug fix (non-breaking change which fixes an issue)
• [x] New feature (non-breaking change which adds functionality)
• [ ] This change requires a documentation update

How has this change been tested, please provide a testcase or example of how you tested the change?

import torch
import requests
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForSegmentation

processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(images=image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_segmentation(
    outputs=outputs, target_sizes=target_size)[0]

detections = sv.Detections.from_transformers(results)
mask_annotator = sv.MaskAnnotator()
annotated_image = mask_annotator.annotate(scene=image, detections=detections)

Any specific deployment considerations

For segmentation models, the processor must return the masks key.

Docs

• [ ] Docs updated? What were the changes:

[CLAassistant]

All committers have signed the CLA.

[SkalskiP]

Hi @Griffin-Sullivan 👋🏻 Thanks for your interest in supervision and time preparing this PR and Colab.

1. Please accept the CLA. Click the yellow button above and follow the instructions.
2. I left a comment with suggested changes that will (most likely) solve your issue.

[Griffin-Sullivan]

Ok I updated everything and it seemed to work fine. Let me know if there's any way we can improve this.

[SkalskiP]

@Griffin-Sullivan This is awesome! I'll update our docs page to include information about segmentation models being supported.