opencv 推理yolov5

opencv yolov5推理

使用opencv推理不带后处理节点的onnx模型。
不带后处理的onnx模型可以用任意宽高的图片作为模型输入。

1. 删除后续节点，这里提供代码和onnx结构，可以根据自己的onnx文件进行修改

import onnx
onnx_file = "yolov5.onnx"
save = "yolov5_del.onnx"
model = onnx.load(onnx_file)

node = model.graph.node

index = []
// 记录要删除的节点
for i in range(len(node)):
    if  node[i].name == "Reshape_317":
        index.append(i)
    if  node[i].name == "Transpose_318":
        print("Transpose_842",i)
        index.append(i)
    
    if  node[i].name == "Reshape_301":
        index.append(i)
    if  node[i].name == "Transpose_302":
        index.append(i)

    if  node[i].name == "Reshape_285":
        index.append(i)
    if  node[i].name == "Transpose_286":
        index.append(i)



print(index)
for i in reversed(index):
    node.remove(node[i])
out = model.graph.output

del out[0]
// 更改输出
out[1].name = "540"
out[1].type.tensor_type.shape.dim[0].dim_value = 1
out[1].type.tensor_type.shape.dim[1].dim_value = 126
out[1].type.tensor_type.shape.dim[2].dim_value = 64
out[1].type.tensor_type.shape.dim[3].dim_value = 112

del out[1].type.tensor_type.shape.dim[4]

out[2].name = "560"
out[2].type.tensor_type.shape.dim[0].dim_value = 1
out[2].type.tensor_type.shape.dim[1].dim_value = 126
out[2].type.tensor_type.shape.dim[2].dim_value = 32
out[2].type.tensor_type.shape.dim[3].dim_value = 56

del out[2].type.tensor_type.shape.dim[4]


out[3].name = "580"
out[3].type.tensor_type.shape.dim[0].dim_value = 1
out[3].type.tensor_type.shape.dim[1].dim_value = 126
out[3].type.tensor_type.shape.dim[2].dim_value = 16
out[3].type.tensor_type.shape.dim[3].dim_value = 28

del out[3].type.tensor_type.shape.dim[4]
onnx.save(model, save)
onnx.checker.check_model(model)

yolov5_onnx.png

3. 推理代码

// letter box
auto  image = imread(path);
float scale_x = m_input_width  / (float)image.cols;
float scale_y = m_input_height / (float)image.rows;
float scale = min(scale_x, scale_y);
float i2d[6], d2i[6];
i2d[0] = scale;  i2d[1] = 0;  i2d[2] = (-scale * image.cols + m_input_width + scale  - 1) * 0.5;
i2d[3] = 0;  i2d[4] = scale;  i2d[5] = (-scale * image.rows + m_input_height + scale - 1) * 0.5;

Mat m2x3_i2d(2, 3, CV_32F, i2d);
Mat m2x3_d2i(2, 3, CV_32F, d2i);
invertAffineTransform(m2x3_i2d, m2x3_d2i);

Mat input_image(m_input_height, m_input_width, CV_8UC3);
warpAffine(image, input_image, m2x3_i2d, input_image.size(), cv::INTER_LINEAR, cv::BORDER_CONSTANT, cv::Scalar::all(114));
// letter box 

Mat blob = blobFromImage(input_image, 1/255.0, Size(m_input_width, m_input_height), Scalar(0, 0, 0), true, false);
m_net.setInput(blob);
vector<Mat> outs;
m_net.forward(outs, m_net.getUnconnectedOutLayersNames());
map<int, vector<float>> confidences;
map<int, vector<Rect>> boxes;
map<int, vector<int>> classIds;
float dethreshold = desigmoid(m_min_conf_threshold);

int n = 0, q = 0, i = 0, j = 0, nout = m_class_names.size() + 5, c = 0, area = 0;
for (n = 0; n < 3; n++)
{
    int num_grid_x = (int)(m_input_width / m_stride[n]);
    int num_grid_y = (int)(m_input_height / m_stride[n]);
    area           = num_grid_x * num_grid_y;

    float* pdata   = (float*)outs[n].data;
    for (q = 0; q < 3; q++) 
    {
        const float anchor_w = m_anchors[n][q * 2];
        const float anchor_h = m_anchors[n][q * 2 + 1];
        for (i = 0; i < num_grid_y; i++)
        {
            for (j = 0; j < num_grid_x; j++)
            {
                float box_score = pdata[4 * area + i * num_grid_x + j];
                if (box_score > dethreshold)
                {
                    float max_class_socre = -100000, class_socre = 0;
                    int max_class_id = 0;
                    for (c = 0; c < m_num_class; c++) //// get max socre
                    {
                        class_socre = pdata[(c + 5) * area + i * num_grid_x + j];
                        if (class_socre > max_class_socre)
                        {
                            max_class_socre = class_socre;
                            max_class_id = c;
                        }
                    }
                    if (max_class_socre < dethreshold) continue;

                    max_class_socre = sigmoid(max_class_socre) * sigmoid(box_score);

                    if (max_class_socre > m_conf_threshold[max_class_id])
                    {
                        float cx   = (sigmoid(pdata[i * num_grid_x + j]) * 2.f - 0.5f + j) * m_stride[n];  ///cx
                        float cy   = (sigmoid(pdata[area + i * num_grid_x + j]) * 2.f - 0.5f + i) * m_stride[n];   ///cy
                        float w    = powf(sigmoid(pdata[2 * area + i * num_grid_x + j]) * 2.f, 2.f) * anchor_w;   ///w
                        float h    = powf(sigmoid(pdata[3 * area + i * num_grid_x + j]) * 2.f, 2.f) * anchor_h;  ///h

                        int left   = int(cx - 0.5*w);
                        int top    = int(cy - 0.5*h);
                        int right  = left + w;
                        int bottom = top + h;
                       
                        affine_project(d2i, left,  top,    &left,  &top);
                        affine_project(d2i, right, bottom, &right, &bottom); 

                        confidences[max_class_id].push_back(max_class_socre);
                        boxes[max_class_id].push_back(Rect(left, top, right-left, bottom-top));
                        classIds[max_class_id].push_back(max_class_id);
                    }
                }
            }
        }
        pdata += area * nout;
    }
}
vector<Box> result_boxes;
for (auto iter : classIds)
{
    vector<int> indices;
    int key = iter.first;
    NMSBoxes(boxes[key], confidences[key], m_min_conf_threshold, m_nms_threshold, indices);
    for (size_t i = 0; i < indices.size(); ++i)
    {
        int idx = indices[i];
        Rect box = boxes[key][idx];
        Box b(m_class_names[key], box.x, box.y, box.x + box.width, box.y + box.height, confidences[key][idx])；
        result_boxes.emplace_back(b);
    }
}

这里使用opencv的NMSBoxes来做后处理，但是opencv的nms是针对所有类别的，所以需要将每一类的识别结果存储起来，然后对每一类做nms，在整合结果。在查阅资料过程中，发现opencv4.6-dev实现了NMSBoxesBatched方法，可以单独对每一类做nms。