nielsr HF staff commited on
Commit
0631208
1 Parent(s): 073311f

Add code example

Browse files
Files changed (1) hide show
  1. README.md +10 -4
README.md CHANGED
@@ -28,16 +28,22 @@ You can use the raw model for object detection. See the [model hub](https://hugg
28
  Here is how to use this model:
29
 
30
  ```python
31
- from transformers import ViTFeatureExtractor, ViTModel
32
  from PIL import Image
33
  import requests
 
34
  url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
35
  image = Image.open(requests.get(url, stream=True).raw)
36
- feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
37
- model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
 
 
38
  inputs = feature_extractor(images=image, return_tensors="pt")
39
  outputs = model(**inputs)
40
- last_hidden_states = outputs.last_hidden_state
 
 
 
41
  ```
42
 
43
  Currently, both the feature extractor and model support PyTorch.
 
28
  Here is how to use this model:
29
 
30
  ```python
31
+ from transformers import DetrFeatureExtractor, DetrForObjectDetection
32
  from PIL import Image
33
  import requests
34
+
35
  url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
36
  image = Image.open(requests.get(url, stream=True).raw)
37
+
38
+ feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50')
39
+ model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
40
+
41
  inputs = feature_extractor(images=image, return_tensors="pt")
42
  outputs = model(**inputs)
43
+
44
+ # model predicts bounding boxes and corresponding COCO classes
45
+ logits = outputs.logits
46
+ bboxes = outputs.pred_boxes
47
  ```
48
 
49
  Currently, both the feature extractor and model support PyTorch.