yolov5输出的解释

Question

我正在做一个面罩检测项目，我用超级解析学/yolov5训练了我的模型。我把经过训练的模型保存为onnx文件，您可以在这里找到模型文件model.onnx。现在，我希望您使用这个model.onnx与opencv来检测实时面罩。训练期间输入的图像大小为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)，谁能帮我解释这个形状，我如何使用这个结果数组来预测我的类，边界框和信心。

Answer 1

我从未使用过纯yolov5模型，但这是yolov5s的输出格式。看起来应该是相似的。

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