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