通过样例YOLOV7推理代码适配YOLOV9模型
收藏回复举报
通过样例YOLOV7推理代码适配YOLOV9模型
发表于2024-08-13 18:26:18
0 查看

1.概述

    基于参考Sample仓官方示例中的sampleYOLOV7进行的YOLOV9适配。需要注意的是,YOLOV7模型输出的数据大小为[1,25200,85],而YOLOV9模型输出的数据大小为[1,84,8400],因此,需要对sampleYOLOV7中的后处理部分进行修改,两个模型的输入一致所以不需要对预处理做更多改变,从而做到YOLOV9模型的适配。

2.准备Yolov9模型

       首先我们从github上获取YOLOV9的模型文件(GitHub - WongKinYiu/yolov9: Implementation of paper - YOLOv9)下载连接中yolov9-c-converted.pt模型及代码,在昇腾设备上安装yolov9环境,并执行github代码仓中的exprot.py脚本将.pt模型转化成onnx模型

#从base环境创建新的环境yolov9
conda create -n yolov9 --clone base
#激活虚拟环境yolov9
conda activate yolov9

安装依赖并执行转换脚本

# 安装yolov9项目的依赖
(yolov9) root@orangepiaipro:~/liwei/yolov9-main/yolov9-main# pip3 install -r requirements.txt
# 模型转换导出onnx
(yolov9) root@orangepiaipro:~/liwei/yolov9-main/yolov9-main# python3 export.py --weights yolov9-c-converted.pt --include onnx

执行之后回显,生成yolov9-c-converted.onnx模型成功

(yolov9) root@orangepiaipro:~/liwei/yolov9-main/yolov9-main# python3 export.py --weights yolov9-c-converted.pt --include onnx
export: data=data/coco.yaml, weights=['yolov9-c-converted.pt'], imgsz=[640, 640], batch_size=1, device=cpu, half=False, inplace=False, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['onnx']
YOLO 🚀 2024-7-25 Python-3.9.2 torch-2.1.0 CPU

Fusing layers...
gelan-c summary: 387 layers, 25288768 parameters, 64944 gradients, 102.1 GFLOPs

PyTorch: starting from yolov9-c-converted.pt with output shape (1, 84, 8400) (49.1 MB)

ONNX: starting export with onnx 1.16.1...
ONNX: export success ✅ 13.8s, saved as yolov9-c-converted.onnx (96.8 MB)

Export complete (24.2s)
Results saved to /root/liwei/yolov9-main/yolov9-main
Detect:          python detect.py --weights yolov9-c-converted.onnx
Validate:        python val.py --weights yolov9-c-converted.onnx
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov9-c-converted.onnx')
Visualize:       https://netron.app

用 netron 工具进行查看转化成功的yolov9-c-converted.onnx模型 

从输出信息中可以看出, 转换后的的模型输出尺寸为 (1, 84, 8400)
转化后的yolov9-c-converted.onnx模型输入的节点名称为images,输入张量的大小为[1,3,640,640],下一步使用昇腾设备香橙派进行ATC离线模型转换:

atc --model=yolov9-c-converted.onnx --framework=5 --output=yolov9n --input_shape="images:1,3,640,640" --soc_version=Ascend310B4 --insert_op_conf=aipp.cfg

--input_shape="images:1,3,640,640"  表示NCHW,即batchsize为1,通道为3,图片大小为640x640
--insert_op_conf=aipp.cfg   
其中aipp做了色域转换为RGB格式,裁剪图片宽高为640,并进行归一化配置,aipp.cfg内容如下:

aipp_op{
    aipp_mode:static
    input_format : YUV420SP_U8
    src_image_size_w : 640
    src_image_size_h : 640
    
    csc_switch : true
    rbuv_swap_switch : false
    matrix_r0c0 : 256
    matrix_r0c1 : 0
    matrix_r0c2 : 359
    matrix_r1c0 : 256
    matrix_r1c1 : -88
    matrix_r1c2 : -183
    matrix_r2c0 : 256
    matrix_r2c1 : 454
    matrix_r2c2 : 0
    input_bias_0 : 0
    input_bias_1 : 128
    input_bias_2 : 128
    
    crop: true
    load_start_pos_h : 0
    load_start_pos_w : 0
    crop_size_w : 640
    crop_size_h : 640

    min_chn_0 : 0
    min_chn_1 : 0
    min_chn_2 : 0
    var_reci_chn_0: 0.0039215686274509803921568627451
    var_reci_chn_1: 0.0039215686274509803921568627451
    var_reci_chn_2: 0.0039215686274509803921568627451
}

 执行成功后,会生成 yolov9n.om 离线模型

3.代码适配

YOLOV7模型输出的数据大小为[1,25200,85]
YOLOV9模型输出的数据大小为[1,84,8400]

  • 第一个维度 1: 代表批次大小。在推理时,通常一次只处理一张图片,因此批次大小为1。
  • 第二个维度 25200,8400: 代表所有预测框的数量。YOLOV7在三个不同的尺度上进行预测,分别用于检测不同大小的目标,大目标中等目标及小目标,输出25200个框; yolov9模型8400个框各尺度输出特征图叠加之后的结果。
  • 第三个维度 85,84: 代表每个预测框的输出特征,包括4个用于边界框位置参数(x, y, w, h),1个置信度得分,置信度是指这个框对包含对象的概率,以及80个类别得分(COCO 数据集的 80 个类别)。其中YOLOV9模型的区别在于检测框的数量不同及不另外对置信度进行预测, 其84 代表 4个位的边界框预测值(x,y,w,h)位置信息和80个检测类别数

由于两个模型的输入一致,所以不需要对预处理进行修改,需要修改后处理代码中检测框的数量以及不再以置信度得分作为其中一个筛选,直接遍历所有框中每个类别的得分找出最高的,YOLOV9的输出的数据可以理解为84行8400列,再由每一行依次变换到首行的末尾构成一维数组,后处理代码如下:

Result SampleYOLOV9::GetResult(std::vector<InferenceOutput> &inferOutputs,
                               string imagePath, size_t imageIndex, bool release)
{
    uint32_t outputDataBufId = 0;
    float *classBuff = static_cast<float *>(inferOutputs[outputDataBufId].data.get());
    // 设置置信度阈值
    float confidenceThreshold = 0.25;
    //支持检测的类别数量
    size_t classNum = 80;
    //设置偏移量x,y,width,height
    size_t offset = 4;
    // 模型输出检测框总数 [1,84,8400]
    size_t modelOutputBoxNum = 8400; 
    // 读取图片
    cv::Mat srcImage = cv::imread(imagePath);
    int srcWidth = srcImage.cols;
    int srcHeight = srcImage.rows;
    // 初始化存储最红过滤后的检测框容器
    vector<BoundBox> boxes;
    size_t yIndex = 1;
    size_t widthIndex = 2;
    size_t heightIndex = 3;
    //循环遍历模型输出框中每个类别的置信度
    for (size_t i = 0; i < modelOutputBoxNum; ++i)
    {
        float maxValue = 0;
        size_t maxIndex = 0;
        for (size_t j = 0; j < classNum; ++j)
        {
            float value = classBuff[(offset + j) * modelOutputBoxNum + i];
            if (value > maxValue)
            {
                maxIndex = j;
                maxValue = value;
            }
        }
     //保存最大置信度的框信息    
        if (maxValue > confidenceThreshold)
        {
            BoundBox box;
            box.x = classBuff[i] * srcWidth / modelWidth_;
            box.y = classBuff[yIndex * modelOutputBoxNum + i] * srcHeight / modelHeight_;
            box.width = classBuff[widthIndex * modelOutputBoxNum + i] * srcWidth / modelWidth_;
            box.height = classBuff[heightIndex * modelOutputBoxNum + i] * srcHeight / modelHeight_;
            box.score = maxValue;
            box.classIndex = maxIndex;
            box.index = i;
            if (maxIndex < classNum)
            {
                boxes.push_back(box);
            }
        }
    }

    ACLLITE_LOG_INFO("filter boxes by confidence threshold > %f success, boxes size is %ld", confidenceThreshold,boxes.size());

    // NMS后处理保留最优框
    vector<BoundBox> result;
    result.clear();
    //设置NMS的阈值,用于判断两个候选框重叠(交并比IOU),IOU超过此阈值的框会被过滤
    float NMSThreshold = 0.45;
    int32_t maxLength = modelWidth_ > modelHeight_ ? modelWidth_ : modelHeight_;
    //按照置信度进行排序
    std::sort(boxes.begin(), boxes.end(), sortScore);
    BoundBox boxMax;
    BoundBox boxCompare;
    //逐个对比并应用NMS保留重叠率较低的框
    while (boxes.size() != 0)
    {
        size_t index = 1;
        result.push_back(boxes[0]);
        while (boxes.size() > index)
        {
            boxMax.score = boxes[0].score;
            boxMax.classIndex = boxes[0].classIndex;
            boxMax.index = boxes[0].index;

            // 不同类别候选框对x,y坐标进行不同的偏移,防止互相干扰
            boxMax.x = boxes[0].x + maxLength * boxes[0].classIndex;
            boxMax.y = boxes[0].y + maxLength * boxes[0].classIndex;
            boxMax.width = boxes[0].width;
            boxMax.height = boxes[0].height;

            boxCompare.score = boxes[index].score;
            boxCompare.classIndex = boxes[index].classIndex;
            boxCompare.index = boxes[index].index;

            //待比较的框做同样的操作偏移
            boxCompare.x = boxes[index].x + boxes[index].classIndex * maxLength;
            boxCompare.y = boxes[index].y + boxes[index].classIndex * maxLength;
            boxCompare.width = boxes[index].width;
            boxCompare.height = boxes[index].height;

            // 计算相交区域的左上和右下坐标及面积
            float xLeft = max(boxMax.x, boxCompare.x);
            float yTop = max(boxMax.y, boxCompare.y);
            float xRight = min(boxMax.x + boxMax.width, boxCompare.x + boxCompare.width);
            float yBottom = min(boxMax.y + boxMax.height, boxCompare.y + boxCompare.height);
            float width = max(0.0f, xRight - xLeft);
            float hight = max(0.0f, yBottom - yTop);
            float area = width * hight;
            //iou = 相交面积/并集面积
            float iou = area / (boxMax.width * boxMax.height + boxCompare.width * boxCompare.height - area);
            //大于阈值则删除
            if (iou > NMSThreshold)
            {
                boxes.erase(boxes.begin() + index);
                continue;
            }
            ++index;
        }
        boxes.erase(boxes.begin());
    }

    ACLLITE_LOG_INFO("filter boxes by NMS threshold > %f success, result size is %ld", NMSThreshold,result.size());
 
    // opencv 画框
    const double fountScale = 0.5;
    const uint32_t lineSolid = 2;
    const uint32_t labelOffset = 11;
    const cv::Scalar fountColor(0, 0, 255); // BGR
    const vector<cv::Scalar> colors{
        cv::Scalar(255, 0, 0), cv::Scalar(0, 255, 0),
        cv::Scalar(0, 0, 255)};

    int half = 2;
    for (size_t i = 0; i < result.size(); ++i)
    {
        cv::Point leftUpPoint, rightBottomPoint;
        leftUpPoint.x = result[i].x - result[i].width / half;
        leftUpPoint.y = result[i].y - result[i].height / half;
        rightBottomPoint.x = result[i].x + result[i].width / half;
        rightBottomPoint.y = result[i].y + result[i].height / half;
        cv::rectangle(srcImage, leftUpPoint, rightBottomPoint, colors[i % colors.size()], lineSolid);
        string className = label[result[i].classIndex];
        string markString = to_string(result[i].score) + ":" + className;

        ACLLITE_LOG_INFO("object detect [%s] success", markString.c_str());

        cv::putText(srcImage, markString, cv::Point(leftUpPoint.x, leftUpPoint.y + labelOffset),
                    cv::FONT_HERSHEY_COMPLEX, fountScale, fountColor);
    }
    string savePath = "out_" + to_string(imageIndex) + ".jpg";
    cv::imwrite(savePath, srcImage);
    if (release)
    {
        free(classBuff);
        classBuff = nullptr;
    }
    return SUCCESS;
}

4 编译运行

bash sample_build.sh
bash sample_run.sh
[INFO] The sample starts to run
[INFO]  Acl init ok
[INFO]  Open device 0 ok
[INFO]  Use default context currently
[INFO]  dvpp init resource ok
[INFO]  Load model ../model/yolov9n.om success
[INFO]  Create model description success
[INFO]  Create model(../model/yolov9n.om) output success
[INFO]  Init model ../model/yolov9n.om success
[INFO]  filter boxes by confidence threshold > 0.350000 success, boxes size is 10
[INFO]  filter boxes by NMS threshold > 0.450000 success, result size is 1
[INFO]  object detect [0.955078:dog] success
[INFO]  Unload model ../model/yolov9n.om success
[INFO]  destroy context ok
[INFO]  Reset device 0 ok
[INFO]  Finalize acl ok
[INFO] The program runs successfully

在out目录下打开检测结果图片

我要发帖子