用opencv推理yolo v5模型,输出带矩形检测框的图片
时间: 2023-05-25 22:06:00 浏览: 224
以下是用OpenCV推理Yolo V5模型并输出带矩形检测框的图片的Python代码实现:
```python
import cv2
import numpy as np
# 加载模型和类别标签
net = cv2.dnn.readNetFromDarknet("yolov5.cfg", "yolov5.weights")
classes = []
with open("coco.names", "r") as f:
classes = [line.strip() for line in f.readlines()]
# 加载图片
img = cv2.imread("test.jpg")
# 获取图片尺寸
height, width, channels = img.shape
# 创建输入blob并执行前向推理
blob = cv2.dnn.blobFromImage(img, 1/255.0, (416, 416), swapRB=True, crop=False)
net.setInput(blob)
outputs = net.forward(net.getUnconnectedOutLayersNames())
# 获取检测框及其置信度和类别
class_ids = []
confidences = []
boxes = []
for output in outputs:
for detection in output:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > 0.5:
center_x = int(detection[0] * width)
center_y = int(detection[1] * height)
w = int(detection[2] * width)
h = int(detection[3] * height)
x = int(center_x - w/2)
y = int(center_y - h/2)
boxes.append([x, y, w, h])
confidences.append(float(confidence))
class_ids.append(class_id)
# 应用非最大抑制算法
indices = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
# 输出带矩形检测框的图片
for i in indices.flatten():
x, y, w, h = boxes[i]
label = str(classes[class_ids[i]])
confidence = str(round(confidences[i], 2))
color = (0, 255, 0)
cv2.rectangle(img, (x, y), (x+w, y+h), color, 2)
cv2.putText(img, label + " " + confidence, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
cv2.imwrite("result.jpg", img)
```
注:在运行以上代码前,请确保已经成功下载了Yolo V5模型文件(yolov5.cfg和yolov5.weights)、类别标签文件(coco.names)以及待检测的图片文件(test.jpg)。另外,需要安装OpenCV库和Numpy库。