读取图片,使用YOLOv5中的训练模型来检测,再输出检测结果图片

以下是使用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方法来绘制结果。