Add code example

This commit is contained in:
Niels Rogge 2021-06-01 10:14:31 +00:00 committed by huggingface-web
parent 073311fe2e
commit 0631208742
1 changed files with 10 additions and 4 deletions

View File

@ -28,16 +28,22 @@ You can use the raw model for object detection. See the [model hub](https://hugg
Here is how to use this model:
```python
from transformers import ViTFeatureExtractor, ViTModel
from transformers import DetrFeatureExtractor, DetrForObjectDetection
from PIL import Image
import requests
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50')
model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
inputs = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
last_hidden_states = outputs.last_hidden_state
# model predicts bounding boxes and corresponding COCO classes
logits = outputs.logits
bboxes = outputs.pred_boxes
```
Currently, both the feature extractor and model support PyTorch.