使用yolov5的onnx模型预测目标的脚本

以下是使用 YOLOv5 的 ONNX 模型进行目标预测的 Python 脚本：

import onnxruntime
import cv2
import numpy as np

# 加载 ONNX 模型
sess = onnxruntime.InferenceSession("yolov5.onnx")

# 读取图像
img = cv2.imread("image.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 对图像进行预处理
img = cv2.resize(img, (640, 640))
img = img.transpose(2, 0, 1)
img = img.astype("float32")
img /= 255.0
img = np.expand_dims(img, axis=0)

# 进行目标检测
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
output = sess.run([output_name], {input_name: img})[0]

# 解析预测结果
boxes = []
scores = []
classes = []
for detection in output:
    x1, y1, x2, y2, score, class_id = detection
    if score > 0.5:
        boxes.append([x1, y1, x2, y2])
        scores.append(score)
        classes.append(class_id)

# 在图像上绘制检测结果
for i in range(len(boxes)):
    x1, y1, x2, y2 = boxes[i]
    label = f"{classes[i]}: {scores[i]:.2f}"
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 显示图像
cv2.imshow("result", img)
cv2.waitKey(0)
cv2.destroyAllWindows()

其中，需要先将 YOLOv5 的 PyTorch 模型转换为 ONNX 模型，可以使用以下命令：

$ python models/export.py --weights yolov5s.pt --img-size 640 --batch-size 1 --include onnx

该命令将在 `models` 目录下生成 `yolov5.onnx` 文件。在预测时，需要将图像先进行预处理，然后通过 `sess.run` 方法进行目标检测，最后解析检测结果并在图像上绘制。