自定义YOLOv5 ONNX模型C++推理目标检测结果异常问题求助

核心问题定位：维度不匹配 #

你的代码中硬编码了YOLOv5官方模型的输出维度85（const int dimensions = 85;），但自定义训练的模型只有2个类别（苹果、香蕉），每个检测框的输出维度应为5（cx,cy,w,h,conf） + 2（类别得分）=7。步长错误会导致读取错误的置信度、框坐标数据，直接引发异常检测结果。

具体修复与排查步骤 #

1. 修正后处理的维度读取逻辑 #

不要硬编码维度，从输出Tensor中动态获取实际维度：

Mat post_process(Mat &input_image, vector<Mat> &outputs, const vector<string> &class_name)
{
    vector<int> class_ids;
    vector<float> confidences;
    vector<Rect> boxes;

    // 动态获取输出维度（替代硬编码的85）
    int rows = outputs[0].size[1]; // 检测框数量，通常为25200
    int dimensions = outputs[0].size[2]; // 自定义模型应为7
    float *data = (float *)outputs[0].data;

    float x_factor = input_image.cols / INPUT_WIDTH;
    float y_factor = input_image.rows / INPUT_HEIGHT;

    for (int i = 0; i < rows; ++i) 
    {
        float confidence = data[4];
        if (confidence >= CONFIDENCE_THRESHOLD) 
        {
            float * classes_scores = data + 5;
            Mat scores(1, class_name.size(), CV_32FC1, classes_scores);
            Point class_id;
            double max_class_score;
            minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
            if (max_class_score > SCORE_THRESHOLD) 
            {
                confidences.push_back(confidence);
                class_ids.push_back(class_id.x);

                float cx = data[0];
                float cy = data[1];
                float w = data[2];
                float h = data[3];
                int left = int((cx - 0.5 * w) * x_factor);
                int top = int((cy - 0.5 * h) * y_factor);
                int width = int(w * x_factor);
                int height = int(h * y_factor);
                boxes.push_back(Rect(left, top, width, height));
            }
        }
        data += dimensions; // 用实际维度步进，替代固定的85
    }

    // NMS与绘图逻辑保持不变
    vector<int> indices;
    NMSBoxes(boxes, confidences, SCORE_THRESHOLD, NMS_THRESHOLD, indices);
    for (int i = 0; i < indices.size(); i++) 
    {
        int idx = indices[i];
        Rect box = boxes[idx];

        int left = box.x;
        int top = box.y;
        int width = box.width;
        int height = box.height;
        rectangle(input_image, Point(left, top), Point(left + width, top + height), BLUE, 3*THICKNESS);

        string label = format("%.2f", confidences[idx]);
        label = class_name[class_ids[idx]] + ":" + label;
        draw_label(input_image, label, left, top);
    }

    return input_image;
}

2. 检查ONNX模型导出参数 #

• 导出时必须指定和C++代码一致的输入尺寸：

  python export.py --weights best.pt --include onnx --imgsz 640 640
  

• 用Netron工具打开ONNX模型，确认输入维度为[1,3,640,640]，输出维度为[1,25200,7]。

3. 对齐图像预处理逻辑 #

YOLOv5训练时默认使用Letterbox缩放（保持图像比例并填充黑边），而你的C++代码直接将图像拉伸到640x640，会导致框坐标映射错误。需修改预处理逻辑：

vector<Mat> pre_process(Mat &input_image, Net &net)
{
    // Letterbox缩放
    Mat letterbox_img;
    float scale = min(INPUT_WIDTH / input_image.cols, INPUT_HEIGHT / input_image.rows);
    int new_w = input_image.cols * scale;
    int new_h = input_image.rows * scale;
    resize(input_image, letterbox_img, Size(new_w, new_h));

    // 填充黑边（YOLOv5默认填充值为114）
    int top = (INPUT_HEIGHT - new_h) / 2;
    int left_pad = (INPUT_WIDTH - new_w) / 2;
    copyMakeBorder(letterbox_img, letterbox_img, top, INPUT_HEIGHT - new_h - top, left_pad, INPUT_WIDTH - new_w - left_pad, BORDER_CONSTANT, Scalar(114,114,114));

    // 转Blob
    Mat blob;
    blobFromImage(letterbox_img, blob, 1./255., Size(INPUT_WIDTH, INPUT_HEIGHT), Scalar(), true, false);

    net.setInput(blob);
    vector<Mat> outputs;
    net.forward(outputs, net.getUnconnectedOutLayersNames());
    return outputs;
}

同时，后处理的坐标映射需调整（减去填充的偏移量）：

// 替换原x_factor、y_factor的计算
float x_factor = input_image.cols / (float)new_w;
float y_factor = input_image.rows / (float)new_h;
// 计算框坐标时减去填充偏移
int left = int((cx - 0.5 * w - left_pad) * x_factor);
int top = int((cy - 0.5 * h - top) * y_factor);
int width = int(w * x_factor);
int height = int(h * y_factor);

4. 验证类别文件一致性 #

确认obj.names中的类别顺序与训练时data.yaml完全一致，比如第一行是apple，第二行是banana，否则会出现类别标签错误。

5. 调整置信度阈值 #

若模型置信度普遍偏低，可临时调低CONFIDENCE_THRESHOLD至0.2，观察检测结果是否正常，再逐步优化阈值。

内容的提问来源于stack exchange，提问作者techgtec