srinivasgs
commited on
Commit
•
199695c
1
Parent(s):
3686e65
updated the How to use section so that the code actually does what the live demo does
Browse filesadded some code to convert the raw outputs (which are not very useful/interpretable) to run it through the image processor and generate a message with text showing confidence and human-readable labels.
the output of this code snippet is now:
```
Detected remote with confidence 0.994 at location [46.96, 72.61, 181.02, 119.73]
Detected remote with confidence 0.975 at location [340.66, 79.19, 372.59, 192.65]
Detected cat with confidence 0.984 at location [12.27, 54.25, 319.42, 470.99]
Detected remote with confidence 0.922 at location [41.66, 71.96, 178.7, 120.33]
Detected cat with confidence 0.914 at location [342.34, 21.48, 638.64, 372.46]
```
which i think is useful and tells you if the model is working
README.md
CHANGED
@@ -35,7 +35,7 @@ You can use the raw model for object detection. See the [model hub](https://hugg
|
|
35 |
Here is how to use this model:
|
36 |
|
37 |
```python
|
38 |
-
from transformers import YolosFeatureExtractor, YolosForObjectDetection
|
39 |
from PIL import Image
|
40 |
import requests
|
41 |
|
@@ -44,6 +44,7 @@ image = Image.open(requests.get(url, stream=True).raw)
|
|
44 |
|
45 |
feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-tiny')
|
46 |
model = YolosForObjectDetection.from_pretrained('hustvl/yolos-tiny')
|
|
|
47 |
|
48 |
inputs = feature_extractor(images=image, return_tensors="pt")
|
49 |
outputs = model(**inputs)
|
@@ -51,6 +52,17 @@ outputs = model(**inputs)
|
|
51 |
# model predicts bounding boxes and corresponding COCO classes
|
52 |
logits = outputs.logits
|
53 |
bboxes = outputs.pred_boxes
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
```
|
55 |
|
56 |
Currently, both the feature extractor and model support PyTorch.
|
|
|
35 |
Here is how to use this model:
|
36 |
|
37 |
```python
|
38 |
+
from transformers import YolosFeatureExtractor, YolosForObjectDetection, AutoImageProcessor
|
39 |
from PIL import Image
|
40 |
import requests
|
41 |
|
|
|
44 |
|
45 |
feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-tiny')
|
46 |
model = YolosForObjectDetection.from_pretrained('hustvl/yolos-tiny')
|
47 |
+
image_processor = AutoImageProcessor.from_pretrained("hustvl/yolos-tiny")
|
48 |
|
49 |
inputs = feature_extractor(images=image, return_tensors="pt")
|
50 |
outputs = model(**inputs)
|
|
|
52 |
# model predicts bounding boxes and corresponding COCO classes
|
53 |
logits = outputs.logits
|
54 |
bboxes = outputs.pred_boxes
|
55 |
+
|
56 |
+
|
57 |
+
# print results
|
58 |
+
target_sizes = torch.tensor([image.size[::-1]])
|
59 |
+
results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
|
60 |
+
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
|
61 |
+
box = [round(i, 2) for i in box.tolist()]
|
62 |
+
print(
|
63 |
+
f"Detected {model.config.id2label[label.item()]} with confidence "
|
64 |
+
f"{round(score.item(), 3)} at location {box}"
|
65 |
+
)
|
66 |
```
|
67 |
|
68 |
Currently, both the feature extractor and model support PyTorch.
|