deep-learning - yolov5输出的解读

Question

我正在做一个 mask 检测项目，我使用 ultralytics/yolov5 训练了我的模型。我将训练好的模型保存为一个 onnx 文件，你可以在这里找到模型文件 model.onnx .现在我希望你使用这个带有 opencv 的 model.onnx 来检测实时 mask 。训练时输入图像大小为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