smarttiger commited on
Commit
ae9dc1b
1 Parent(s): 4ed0e52

添加hustvl/yolos-tiny 模型 更改模型的路径

Browse files
models/{facebook → ObjectDetection/facebook}/detr-resnet-50.py RENAMED
File without changes
models/ObjectDetection/hustvl/yolos-tiny.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import YolosImageProcessor, YolosForObjectDetection
2
+ from PIL import Image
3
+ import torch
4
+ import requests
5
+
6
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
7
+ image = Image.open(requests.get(url, stream=True).raw)
8
+
9
+ model = YolosForObjectDetection.from_pretrained('hustvl/yolos-tiny', local_files_only=True)
10
+ image_processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny", local_files_only=True)
11
+
12
+ inputs = image_processor(images=image, return_tensors="pt")
13
+ outputs = model(**inputs)
14
+
15
+ # model predicts bounding boxes and corresponding COCO classes
16
+ logits = outputs.logits
17
+ bboxes = outputs.pred_boxes
18
+
19
+
20
+ # print results
21
+ target_sizes = torch.tensor([image.size[::-1]])
22
+ results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
23
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
24
+ box = [round(i, 2) for i in box.tolist()]
25
+ print(
26
+ f"Detected {model.config.id2label[label.item()]} with confidence "
27
+ f"{round(score.item(), 3)} at location {box}"
28
+ )