提供一个yolov5的示例
时间: 2023-10-12 13:15:59 浏览: 34
这是一个使用yolov5模型进行目标检测的示例代码:
```
import torch
import cv2
from PIL import Image
import numpy as np
# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
# 定义类别名称
classes = ['person', 'car', 'truck', 'bus']
# 定义颜色
colors = [[0, 0, 255], [0, 255, 0], [255, 0, 0], [0, 255, 255]]
# 加载图像
img = cv2.imread('test.jpg')
# 将图像转换为PIL格式
img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
# 进行目标检测
results = model(img, size=640)
# 获取检测结果
boxes = results.xyxy[0].numpy()
scores = results.xyxy[0][:, 4].numpy()
classes_idx = results.xyxy[0][:, 5].numpy()
# 绘制边框和类别名称
for box, score, class_idx in zip(boxes, scores, classes_idx):
if score > 0.5:
box = box.astype(np.int32)
class_name = classes[int(class_idx)]
color = colors[int(class_idx)]
cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2)
cv2.putText(img, class_name, (box[0], box[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
# 显示图像
cv2.imshow('image', cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
cv2.waitKey(0)
cv2.destroyAllWindows()
```
在这个示例中,我们首先使用torch.hub加载了yolov5s模型。然后,我们定义了目标类别名称和颜色,加载了一个测试图像,并将其转换为PIL格式。接着,我们使用模型进行目标检测,并获取检测结果。最后,我们遍历检测结果,绘制边框和类别名称,并将结果显示出来。
相关推荐
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![-](https://csdnimg.cn/download_wenku/file_type_column_c1.png)
![-](https://csdnimg.cn/download_wenku/file_type_column_c1.png)
![-](https://csdnimg.cn/download_wenku/file_type_column_c1.png)
![-](https://csdnimg.cn/download_wenku/file_type_column_c1.png)
![-](https://csdnimg.cn/download_wenku/file_type_column_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)