yolov5解读output

Question

I am making a face mask detection project and I trained my model using ultralytics/yolov5.I saved the trained model as an onnx file, you can find the model file here model.onnx . 现在我想让你使用这个 model.onnx 和 opencv 来检测实时面罩。 训练时的输入图像大小为 320*320。 您可以使用 netron 可视化此 model。 我已经编写了这段代码来使用网络摄像头捕获图像并将其传递给 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 的 output 是 (1, 1, 3, 40, 40, 85) 谁能帮我解释这个形状，我如何使用这个结果数组来预测我的 class、边界框和置信度。

Answer 1

我从未使用过纯 yolov5 model，但这里是 yolov5s 的 output 格式。 它看起来应该是相似的。

ouput tensor structure (yolov5s):
output_tensor[a, b, c, d]
    a -> image index (If you're input is a batch of images, this tells you which image's output you're looking at. If your input is just one image, leave this as 0.)
    b -> index of image in batch
    c -> information about bounding box
        0, 1 -> x and y coordinate of bounding box center
        2, 3 -> width and height of bounding box
        4 -> bounding box confidence
        5 - 85 -> single class confidences
    d -> index of proposed bounding boxes