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提供的导出文件后就正常了。
