Ubuntu+tensorRT算法开发实战

服务器环境：Ubuntu16.04、CUDA11.1、CUDNN8.0、OPENCV4.1.1、TENSORRT7.2

该项目基于极视角cvmart提供的打榜平台，主要内容为口罩佩戴情况检测。

本文仅为个人记录，请谨慎参考

环境配置

算法开发的环境配置不需要自己操作，c++相关的环境都配置好了，cmake等可以直接使用。

由于需要将Pytorch模型转换为onnx模型，模型开发的环境需要增加onnx依赖包，执行如下命令就好了。

pip install onnx==1.9.0
pip install onnx-simplifier==0.3.6

下载封装代码demo

我用的是行人闯入检测的，也可以用车牌车辆识别。

极市地址：https://gitee.com/cvmart/ev_sdk_demo4.0_pedestrian_cid_yolov5/repository/blazearchive/master.zip?Expires=1667485765&Signature=WLUp1ZfbvhCtrNtGJ6TGxcvOk6BeDNgyAblC72JFQ2c%3D

模型转换

需要注意的是，针对采用yolov5框架训练的模型转换，cvmart提供的demo代码里有专门的导出文件，请使用3rd/export.py进行导出，避免后续模型精度出现问题。这一步请在模型开发的环境下完成，执行训练时运行bash

python export.py --data data/cvmart.yaml --weights /project/train/models/exp3/weights/best.pt --simplify --include onnx
# weights中填你的模型文件

修改代码

配置文件

主要就是修改一些标签类别。

• config/algo_config.json

"mark_text_en":["front_wear","front_no_wear","front_under_nose_wear","front_under_mouth_wear","mask_front_wear","mask_front_under_nose_wear","mask_front_under_mouth_wear","side_wear","side_no_wear","side_under_nose_wear","side_under_mouth_wear","mask_side_wear","mask_side_under_nose_wear","mask_side_under_mouth_wear","side_back_head_wear","side_back_head_no_wear","back_head","front_unknown","side_unknown","strap"],
"mark_text_zh": ["正面佩戴口罩包住鼻子","正面未佩戴口罩","正面佩戴口罩在鼻子下且在嘴巴上面","正面佩戴口罩在嘴巴下面","不能很明确知道正面是否佩戴口罩","侧面佩戴口罩包住鼻子","侧面未佩戴口罩","侧面佩戴口罩在鼻子下且在嘴巴上面","侧面佩戴口罩在嘴巴下面","不能很明确知道侧面是否佩戴口罩","侧后方人头带了口罩","侧后方人头没有戴口罩","标注背面人头","标注正面口罩，口罩包住鼻子","标注正面口罩，口罩在鼻子下且在嘴巴上面","标注正面口罩，口罩在嘴巴下面","标注侧面口罩，口罩包住鼻子","标注侧面口罩，口罩在鼻子下且在嘴巴上面","标注侧面口罩，口罩在嘴巴下面","口罩的带子"]

• src/Configuration.hpp

std::map< std::string, std::vector <std::string> > targetRectTextMap = { {"en",{"same_as_above"}}, {"zh", {"和上面一样"}}};// 检测目标框顶部文字

• src/SampleAlgorithm.cpp，按照你标签的顺序去设置报警类型

std::vector<int> alarmType = {1,2,3,8,9,10};

模型路径

• src/SampleAlgorithm.cpp，我是把原来64位的onnx又转成32位了，主要是对warning强迫症。网上都有转换方法，我就随便找了一个。要注意算法开发的模型挂载路径和模型开发是不一样的。

mDetector->Init("/usr/local/ev_sdk/model/exp3/weights/best32.onnx",mConfig.algoConfig.thresh);

检测器和算法逻辑

• src/SampleDetector.cpp，我是修改了NMS的一些操作，把包含在大检测框内的小框都去掉了。

bool SampleDetector::ProcessImage(const cv::Mat& img, std::vector<BoxInfo>& DetObjs, float thresh)
{
	mThresh = thresh;
    // 省略部分代码
    runNms(DetObjs, 0.5); // 此处NMS的Thresh要和你的模型一致，也可以写一个外部参数传入，这里直接修改
}
void SampleDetector::runNms(std::vector<BoxInfo>& objects, float thresh) 
{
    auto cmp_lammda = [](const BoxInfo& b1, const BoxInfo& b2){return b1.score < b2.score;};
    std::sort(objects.begin(), objects.end(), cmp_lammda);
    for(int i = 0; i < objects.size(); ++i)
    {
        if( objects[i].score < 0.25 )
        {
            continue;
        }
        for(int j = i + 1; j < objects.size(); ++j)
        {
            cv::Rect rect1 = cv::Rect{objects[i].x1, objects[i].y1, objects[i].x2 - objects[i].x1, objects[i].y2 - objects[i].y1};
            cv::Rect rect2 = cv::Rect{objects[j].x1, objects[j].y1, objects[j].x2 - objects[i].x1, objects[j].y2 - objects[j].y1};
            if (IOU(rect1, rect2) > thresh)  
            {
                objects[i].score = 0.f;
            }
        }
    }
    auto iter = objects.begin();
    while( iter != objects.end() )
    {
        if(iter->score < 0.25)
        {
            iter = objects.erase(iter);
        }
        else
        {
            ++iter;
        }
    }
}

• src/SampleAlgorithm.cpp

STATUS SampleAlgorithm::Process(const cv::Mat &inFrame, const char *args, JiEvent &event)
{
    // 省略部分代码
    // 算法处理
    cv::Mat img = inFrame.clone();
    mDetector->ProcessImage(img, detectedObjects, mConfig.algoConfig.thresh);   
    std::vector<int> alarmType = {1,2,3,8,9,10};    // 报警类别
    // 过滤出没带口罩的目标添加到validTargets
    for (auto &obj : detectedObjects)
    {
        auto iter = find(alarmType.begin(), alarmType.end(), obj.label);
        if (iter != alarmType.end())
        {
            for (auto &roiPolygon : mConfig.currentROIOrigPolygons)
            {
                int mid_x = (obj.x1 + obj.x2) / 2;
                int mid_y = (obj.y1 + obj.y2) / 2;
                // 当检测的目标的中心点在ROI内的话，就视为闯入ROI的有效目标
                if (WKTParser::inPolygon(roiPolygon, cv::Point(mid_x, mid_y)))
                {
                    validTargets.emplace_back(obj);
                    
                }           
            }
        }
        else{
            continue;
        }
    }
    SDKLOG(INFO) << "detected targets : " << detectedObjects.size() << " valid targets :  " << validTargets.size();
    // 此处示例业务逻辑：当算法检测到有人口罩未戴好时，就报警
    bool isNeedAlert = false; // 是否需要报警
    // 创建输出图
    inFrame.copyTo(mOutputFrame);
    // 画ROI区域
    if (mConfig.drawROIArea && !mConfig.currentROIOrigPolygons.empty())
    {
        drawPolygon(mOutputFrame, mConfig.currentROIOrigPolygons, cv::Scalar(mConfig.roiColor[0], mConfig.roiColor[1], mConfig.roiColor[2]),
                    mConfig.roiColor[3], cv::LINE_AA, mConfig.roiLineThickness, mConfig.roiFill);
    }
    // 判断是否要要报警
    if (validTargets.size() > 0)
    {        
        isNeedAlert = true;
    }
    // 省略后续代码
}

编译测试

编译

直接用cmake去编译就好了。

cd /usr/local/ev_sdk/build
cmake ..
make install
cd /usr/local/ev_sdk/test/build
cmake ..
make install

调试/测试

• 调试就写bash语句，记得把你的数据放在/project/inputs文件夹，图片、视频都可以。指定输出到/project/outputs，这样可以预览文件。

/usr/local/ev_sdk/bin/test-ji-api -f 1 -i /project/ev_sdk/data/demo_mask.jpg -o /project/outputs/result.jpg

• 测试直接用平台发起算法测试就好了。我的模型性能分比较低，不知道是不是我的c++算法太慢了，f-score在0.83左右。对侧脸没带口罩识别率比较低。

小结

基本上用给的demo进行封装很快就可以上手，对模型性能提升还要进一步研究。感谢提供平台以及一些开发人员的指导，让大家能够很顺利地开发并封装，希望能顺利度过新手期，向大神靠拢。

可以看到之前的模型f-score很低，应该是导出时模型精度受到了影响，使用cvmart提供的导出文件后就正常了。