在进行边缘部署的时候,python语言会带来一些不便,如硬件兼容,环境配置等问题,这时候需要用C++语言结合交叉编译,实现智能算法的应用。本篇主要讲推理,训练是使用的python结合pytorch进行的,也就是使用的yolov5官方v7.0文件进行的分割任务训练,转化模型为onnx。
代码地址:GitHub - liushuailiang/yolov5_seg_opencv
运行
代码是结合opencv实现的,首先需要安装opencv库,具体安装方法可以看我上一篇,需要注意的是需要安装里面的dnn模块来做深度学习。 编译安装opencv,C++库,并安装DNN模块_c++ opencv dnn 安装-CSDN博客
安装完成后,准备自己的模型权重(需要是onnx格式,若不是需要该代码中的model/yolo.cpp/25行,改为读取pt文件。而且代码中不涉及模型的构建,所以在用YOLOv5的python源代码训练的时候,最后要保存模型而不是保存模型的参数,也就是torch.save(model, 'net.pth')而不是torch.save(model.state_dict(), 'net_params.pth')),随后准备class.names文件,内容是所有类别的名字,换行符分割。随后修改保存输出文件的文件夹地址,并指明要推理的source,然后可以运行。
cmake .
make
主要用到的参数在main.cpp文件的main函数中可以进行修改。
原理简单叙述
代码是结合opencv库,实现了简单的yolov5多目标分割,并绘制检测结果。首先介绍一下算法原理:
模型输出的是两个Mat,也就是一个vector<Mat>,一个pred代表检测框及协方差系数,尺寸为[1, 25200, 32+5+类别数],一个proto代表特征,尺寸为[1, 32, 160, 160]。25200代表的是yolov5的锚框数量,32是协方差系数的数量,5代表的是框坐标和目标置信度,每个类别对应一个类别分数。以下是模型的输入预处理和输出预处理代码
int col = frame.cols;
int row = frame.rows;
Mat netInputImg;
Vec4d params;
LetterBox(frame, netInputImg, params, cv::Size(inpWidth, inpHeight));
Mat blob = blobFromImage(netInputImg, 1 / 255.0, Size(this->inpWidth, this->inpHeight), Scalar(0, 0, 0), true, false);
this->net.setInput(blob);
vector<Mat> net_output_img;
vector<string> output_layer_names{ "output0","output1" };
net.forward(net_output_img, output_layer_names); //获取output的输出
LetterBox函数用于将输入的图像调整为指定的尺寸(640, 640),同时保持图像的宽高比并填充边界。
下面是后处理。主要针对输出的pred和proto进行处理,得到框坐标、置信度、类别等。
std::vector<int> class_ids;//结果id数组
std::vector<float> confidences;//结果每个id对应置信度数组
std::vector<cv::Rect> boxes;//每个id矩形框
std::vector<vector<float>> picked_proposals; //output0[:,:, 5 + _className.size():net_width]===> for mask
int net_width = net_output_img[0].size[2];
int net_height = net_output_img[0].size[1];
int score_length = net_width - 37;
float* pdata = (float*)net_output_img[0].data;
for (int r = 0; r < net_height; r++) { //lines
float box_score = pdata[4];
if (box_score >= objThreshold) {
// 取出当前锚框的类别得分
cv::Mat scores(1, score_length, CV_32FC1, pdata + 5);
Point classIdPoint;
double max_class_socre;
float cl1 = scores.at<float>(0, 0), cl2 = scores.at<float>(0, 1);
// 获取类别得分的最大值和对应的索引。函数用于找到一个多维数组中的最小值和最大值,以及它们的位置,最小值及其索引不需要,所以设置为0
minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
// 得到置信度最大的类别的置信度
max_class_socre = (float)max_class_socre;
if (max_class_socre >= objThreshold) {
// 根据网络输出提取类别的候选提案信息
vector<float> temp_proto(pdata + 5 + score_length, pdata + net_width);
picked_proposals.push_back(temp_proto);
float x = (pdata[0] - params[2]) / params[0]; //x
float y = (pdata[1] - params[3]) / params[1]; //y
float w = pdata[2] / params[0]; //w
float h = pdata[3] / params[1]; //h
int left = MAX(int(x - 0.5 * w + 0.5), 0);
int top = MAX(int(y - 0.5 * h + 0.5), 0);
class_ids.push_back(classIdPoint.x);
confidences.push_back(max_class_socre * box_score);
boxes.push_back(Rect(left, top, int(w + 0.5), int(h + 0.5)));
}
}
pdata += net_width;//下一行
}
NMS算法实现,通过调用opencv::dnn模块中的NMS,这个NMS不区分类别
vector<int> nms_result;
cv::dnn::NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, nms_result);
std::vector<vector<float>> temp_mask_proposals;
Rect holeImgRect(0, 0, frame.cols, frame.rows);
for (int i = 0; i < nms_result.size(); ++i) {
int idx = nms_result[i];
OutputSeg result;
result.id = class_ids[idx];
result.confidence = confidences[idx];
result.box = boxes[idx] & holeImgRect;
temp_mask_proposals.push_back(picked_proposals[idx]);
output.push_back(result);
}
生成掩码,原理是矩阵乘法,也就是将pred中的协方差系数与proto中的特征进行矩阵相乘,[25200, 32] * [32, 160*160],结果转化为[25200, 160, 160]即为每个锚框得到的分割掩码。
MaskParams mask_params;
mask_params.params = params;
mask_params.srcImgShape = frame.size();
mask_params.maskThreshold = 0.5;
mask_params.netHeight = inpWidth;
mask_params.netWidth = inpWidth;
for (int i = 0; i < temp_mask_proposals.size(); ++i) {
// 根据输入的分割提议(temp_mask_proposals)和分割原型(net_output_img[1]),以及一些参数,生成一个掩码(output.boxMask)。
GetMask2(Mat(temp_mask_proposals[i]).t(), net_output_img[1], output[i], mask_params);
}
绘制结果
void DrawPred(Mat& img, vector<OutputSeg> result, std::vector<std::string> classNames, vector<Scalar> color, bool isVideo) {
Mat mask = img.clone();
for (int i = 0; i < result.size(); i++) {
int left, top;
left = result[i].box.x;
top = result[i].box.y;
int color_num = i;
// 绘制框
rectangle(img, result[i].box, color[result[i].id], 2, 8);
if (result[i].boxMask.rows && result[i].boxMask.cols > 0)
// 把boxMask区域置为某一种随机颜色
mask(result[i].box).setTo(color[result[i].id], result[i].boxMask);
// 绘制标签信息
string label = classNames[result[i].id] + ":" + to_string(result[i].confidence);
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
top = max(top, labelSize.height);
putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
}
// 将原图和绘制了框、Mask、标签的图片按0.5权重结合在一起
addWeighted(img, 0.5, mask, 0.5, 0, img); //add mask to src
}
代码只是基础实现,可能存在问题,但应该都是能改的小问题,还望指正。