YOLOv5语义分割 C++实现

        在进行边缘部署的时候,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
}

代码只是基础实现,可能存在问题,但应该都是能改的小问题,还望指正。

  • 17
    点赞
  • 13
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值