使用onnx预测目标的脚本

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

import onnxruntime
import numpy as np
import cv2

# 加载ONNX模型
session = onnxruntime.InferenceSession("model.onnx")

# 输入图像路径和目标类别列表
def predict(image_path, class_names):
    # 加载并预处理图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (224, 224))
    image = image.astype(np.float32)
    image /= 255.0
    image = np.transpose(image, [2, 0, 1])
    image = np.expand_dims(image, axis=0)

    # 运行预测
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    outputs = session.run([output_name], {input_name: image})

    # 获取预测结果
    prediction = np.argmax(outputs[0][0])
    class_name = class_names[prediction]

    return class_name

# 测试预测函数
class_names = ["cat", "dog"]
image_path = "test.jpg"
class_name = predict(image_path, class_names)
print("Predicted class:", class_name)

在此示例中，我们假设ONNX模型接受大小为224x224的RGB图像作为输入，并返回一个包含目标类别概率的向量。我们首先加载模型，然后定义一个`predict`函数，该函数将图像路径和目标类别列表作为输入，并返回预测的类别名称。在`predict`函数中，我们加载并预处理图像，然后使用ONNX运行预测。最后，我们使用`argmax`函数获取预测结果，并将其映射回目标类别名称。