以下是使用YOLOv5来检测图像并输出结果图像的Python代码示例:
import torch
import cv2
from PIL import Image
# 加载YOLOv5预训练模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
# 读取图像
img = cv2.imread('image.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 转换图像通道顺序
# 检测物体
results = model(Image.fromarray(img))
# 在原图像上绘制检测结果
results.render()
output_img = results.imgs[0]
# 保存输出图像
cv2.imwrite('output.jpg', output_img[:, :, ::-1]) # 再次转换通道顺序
在上面的代码中,我们首先使用torch.hub.load()函数加载了一个预训练的YOLOv5模型。然后,我们使用OpenCV库读取了一张图像,并将其转换为PIL图像格式,以便可以在YOLOv5模型上进行预测。接下来,我们使用模型的results方法来绘制结果。