1

我正在制作一个口罩检测项目,我使用 ultralytics/yolov5 训练了我的模型。我将训练后的模型保存为 onnx 文件,您可以在此处找到模型文件model.onnx。现在我希望你使用这个带有 opencv 的 model.onnx 来检测实时面罩。训练时的输入图像大小为 320*320。您可以使用 netron 可视化此模型。我编写了这段代码来使用网络摄像头捕获图像并将其传递给 model.onnx 以预测我的边界框。代码如下:

def predict(img):
    session = onnxruntime.InferenceSession(model_path)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    img = img.reshape((1,3,320,320))
    data = json.dumps({'data':img.tolist()})
    data = np.array(json.loads(data)['data']).astype('float32')
    result = session.run([output_name],{input_name:data})
    result = np.array(result)
    print(result.shape)

result.shape 的输出是 (1, 1, 3, 40, 40, 85) 谁能帮我解释这个形状,我如何使用这个结果数组来预测我的类、边界框和置信度。

4

0 回答 0