自己实现一个关键点检测模型，并实现对应的c++推理

已于 2024-08-19 13:04:38 修改
阅读量347 收藏 5
点赞数 3
文章标签： c++ 开发语言
于 2024-08-19 11:53:39 首次发布
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/sdhdsf132452/article/details/141319576
版权

文章目录

说明

在这篇文章中，我们使用yolov8实现了关键点的c++推理，在训练层面，是直接调用yolov8做的。为了深入理解关键点训练的过程，这次自己实现一个模型，并将模型转为onnx，实现对应的c++推理。

一、python训练模型实现

记录在这个gitee仓库中，点击访问。但是由于我的手机号不能绑定这个gitee，所以没办法设置为public项目。如果是学生需要的话，可以私信我，我下载下来发给你。

二、c++推理

由于在训练中，我同时实现了

  1. 直接回归关键点
  2. 生成热力图

2种dataset和训练方式，所以c++推理也对应有2种方式

2.1 直接回归关键点的c++推理

#include<iostream>
#include<opencv.hpp>
#include<opencv2/dnn.hpp>
#include<opencv2/imgproc.hpp>

/// <summary>
/// 不失目标比例的缩放图片
/// </summary>
/// <param name="img">输入图片</param>
/// <param name="input_size">模型输入大小[h,w]</param>
/// <param name="out_img">缩放后的图像</param>
/// <returns>缩放成功或失败</returns>
bool keepRatioResizeV2(const cv::Mat& img, const int input_size[2], cv::Mat& out_img)
{
	int rows = img.rows;
	int cols = img.cols;
	double scale = 0.0;

#pragma region 先把较短的边缩放为目标大小
	if (rows < cols)
	{
		scale = input_size[0] * 1.0 / rows;
	}
	else scale = input_size[0] * 1.0 / cols;

	cv::Mat proImg;
	cv::resize(img, proImg, cv::Size(0,0), scale, scale);
	// std::cout << proImg.size << std::endl;// [h,w]
#if 0
	cv::imshow("proImg", proImg);
	cv::waitKey();
#endif
#pragma endregion
	int proImgCols = proImg.cols;
	int proImgRows = proImg.rows;
	if (proImgRows > input_size[0])
	{
		// 再下面裁剪
		cv::Rect roi(0, 0, input_size[0], input_size[0]);
		out_img = proImg(roi);
	}
	else
	{
		// 在左右边裁剪
		int cw = proImgCols - input_size[1];
		cv::Rect roi(cw/2, 0, input_size[0], input_size[0]);
		out_img = proImg(roi);
	}
#if 0
	cv::imshow("out_img", out_img);
	cv::waitKey();
#endif
	if (out_img.cols != input_size[1] || out_img.rows != input_size[0])
	{
		return false;
	}
	return true;
}

/// <summary>
/// 将模型输出点坐标还原到原图
/// </summary>
/// <param name="org_size">原始图像大小[h,w]</param>
/// <param name="input_size">输入模型的图像大小[h,w]</param>
/// <param name="input_pt">模型输出点[x,y]</param>
/// <param name="out_pt">还原到原图上的点[x,y]</param>
/// <returns>还原成功或失败</returns>
bool keepRatioResizeBackV2(
	const int org_size[2], 
	const int input_size[2], 
	const cv::Point& input_pt,
	cv::Point& out_pt)
{
	int orgH = org_size[0];
	int orgW = org_size[1];
	int inputH = input_size[0];
	int inputW = input_size[1];

	double scale = 0.0;
	if (orgH < orgW)
	{
		scale = orgH * 1.0 / inputH;
		// 点位置乘以这个缩放大小
		out_pt.x = input_pt.x * scale;
		out_pt.y = input_pt.y * scale;
		//
		double cw = orgW - input_size[1] * scale;
		out_pt.x += cw / 2;
	}
	else
	{// 由于这种情况是直接裁掉下面的,所以只需要缩放就可以
		scale = orgW * 1.0 / inputW;
		// 点位置乘以这个缩放大小
		out_pt.x = input_pt.x * scale;
		out_pt.y = input_pt.y * scale;
	}

	return true;
}


int main()
{
	cv::dnn::Net model = cv::dnn::readNetFromONNX("testnet.onnx");
	std::cout << "load model complete." << std::endl;

	const int inputSize[2]{ 640,640 };//[h,w]
	const int numPoints = 4;
	cv::Mat img = cv::imread("003.png");
	const int orgSize[2]{ img.rows, img.cols }; //[h,w]
#if 0
	cv::imshow("w", img);
	cv::waitKey();
#endif
	cv::Mat resizeImg;
	bool resizeFlag = keepRatioResizeV2(img, inputSize, resizeImg);
	if (resizeFlag == false)
	{
		std::cout << "resize error." << std::endl;
		return -1;
	}

	// 图像预处理
	cv::Mat inputBlob = cv::dnn::blobFromImage(
		resizeImg, 1.0 / 255.0,
		cv::Size(inputSize[1], inputSize[0]),
		cv::Scalar()
	);
	std::cout << inputBlob.size << std::endl;

	model.setInput(inputBlob, "input");

	//
	cv::Mat outputs = model.forward();
	std::cout << outputs.size << std::endl;
	//std::cout << outputs.channels() << std::endl;
	//
	cv::Mat outputs_1 = outputs.reshape(0, numPoints);
	std::cout << outputs_1.size << std::endl;
	if (outputs_1.rows != numPoints)
	{
		std::cout << "reshape false" << std::endl;
		return -2;
	}

	std::cout << outputs_1 << std::endl;
	int type = outputs_1.type();
	std::cout << "type=" << type << std::endl;

	for (int row = 0; row < numPoints; ++row)
	{
		// 取出每一行
		float* rowPtr = outputs_1.ptr<float>(row);
		int x = rowPtr[0] * inputSize[1];
		int y = rowPtr[1] * inputSize[0];
		
		cv::Point curPt(x, y);
#if 0 在resize图上
		cv::circle(
			resizeImg,
			cv::Point(x, y),
			2,
			cv::Scalar(0, 0, 255),
			-1);
#endif

#if 1 还原回原图
		cv::Point outPt;
		keepRatioResizeBackV2(orgSize, inputSize, curPt, outPt);
#endif
		cv::circle(
			img,
			outPt,
			2,
			cv::Scalar(0, 0, 255),
			-1);

	}
#if 0 在resize图上
	cv::imshow("result", resizeImg);
	cv::waitKey();
#endif

#if 1 在原图上
	cv::imshow("result", img);
	cv::waitKey();
#endif


	return 0;
}

2.2 使用热力图的推理

#include<iostream>
#include<opencv.hpp>
#include<opencv2/dnn.hpp>
#include<opencv2/imgproc.hpp>

/// <summary>
/// 不失目标比例的缩放图片
/// </summary>
/// <param name="img">输入图片</param>
/// <param name="input_size">模型输入大小[h,w]</param>
/// <param name="out_img">缩放后的图像</param>
/// <returns>缩放成功或失败</returns>
bool keepRatioResizeV2(const cv::Mat& img, const int input_size[2], cv::Mat& out_img)
{
	int rows = img.rows;
	int cols = img.cols;
	double scale = 0.0;

#pragma region 先把较短的边缩放为目标大小
	if (rows < cols)
	{
		scale = input_size[0] * 1.0 / rows;
	}
	else scale = input_size[0] * 1.0 / cols;

	cv::Mat proImg;
	cv::resize(img, proImg, cv::Size(0, 0), scale, scale);
	// std::cout << proImg.size << std::endl;// [h,w]
#if 0
	cv::imshow("proImg", proImg);
	cv::waitKey();
#endif
#pragma endregion
	int proImgCols = proImg.cols;
	int proImgRows = proImg.rows;
	if (proImgRows > input_size[0])
	{
		// 再下面裁剪
		cv::Rect roi(0, 0, input_size[0], input_size[0]);
		out_img = proImg(roi);
	}
	else
	{
		// 在左右边裁剪
		int cw = proImgCols - input_size[1];
		cv::Rect roi(cw / 2, 0, input_size[0], input_size[0]);
		out_img = proImg(roi);
	}
#if 0
	cv::imshow("out_img", out_img);
	cv::waitKey();
#endif
	if (out_img.cols != input_size[1] || out_img.rows != input_size[0])
	{
		return false;
	}
	return true;
}

/// <summary>
/// 将模型输出点坐标还原到原图
/// </summary>
/// <param name="org_size">原始图像大小[h,w]</param>
/// <param name="input_size">输入模型的图像大小[h,w]</param>
/// <param name="input_pt">模型输出点[x,y]</param>
/// <param name="out_pt">还原到原图上的点[x,y]</param>
/// <returns>还原成功或失败</returns>
bool keepRatioResizeBackV2(
	const int org_size[2],
	const int input_size[2],
	const cv::Point& input_pt,
	cv::Point& out_pt)
{
	int orgH = org_size[0];
	int orgW = org_size[1];
	int inputH = input_size[0];
	int inputW = input_size[1];

	double scale = 0.0;
	if (orgH < orgW)
	{
		scale = orgH * 1.0 / inputH;
		// 点位置乘以这个缩放大小
		out_pt.x = input_pt.x * scale;
		out_pt.y = input_pt.y * scale;
		//
		double cw = orgW - input_size[1] * scale;
		out_pt.x += cw / 2;
	}
	else
	{// 由于这种情况是直接裁掉下面的,所以只需要缩放就可以
		scale = orgW * 1.0 / inputW;
		// 点位置乘以这个缩放大小
		out_pt.x = input_pt.x * scale;
		out_pt.y = input_pt.y * scale;
	}

	return true;
}

#include <opencv2/opencv.hpp>  

// 定义一个结构体来存储颜色  
struct Colors {
	cv::Scalar red;
	cv::Scalar green;
	cv::Scalar blue;
	cv::Scalar yellow;
	cv::Scalar magenta;
	cv::Scalar cyan;
	cv::Scalar white;
	cv::Scalar black;

	// 构造函数，可以在这里初始化颜色  
	Colors() :
		red(0, 0, 255),       // BGR  
		green(0, 255, 0),
		blue(255, 0, 0),
		yellow(0, 255, 255),
		magenta(255, 0, 255),
		cyan(255, 255, 0),
		white(255, 255, 255),
		black(0, 0, 0)
	{}
};



int main()
{
	cv::dnn::Net model = cv::dnn::readNetFromONNX("testnetMap.onnx");
	std::cout << "load model complete." << std::endl;

	const int inputSize[2]{ 640,640 };//[h,w]
	const int numPoints = 4;
	Colors myColors;
	std::vector<cv::Scalar> colors{ myColors.red, myColors.blue,myColors.cyan,myColors.green };

	cv::Mat img = cv::imread("012.png");
	const int orgSize[2]{ img.rows, img.cols }; //[h,w]
#if 0
	cv::imshow("w", img);
	cv::waitKey();
#endif

	cv::Mat resizeImg;
	bool resizeFlag = keepRatioResizeV2(img, inputSize, resizeImg);
	if (resizeFlag == false)
	{
		std::cout << "resize error." << std::endl;
		return -1;
	}

	// 图像预处理
	cv::Mat inputBlob = cv::dnn::blobFromImage(
		resizeImg, 1.0 / 255.0,
		cv::Size(inputSize[1], inputSize[0]),
		cv::Scalar()
	);
	std::cout << "input_size = " << inputBlob.size << std::endl;

	model.setInput(inputBlob, "input"); 

	//
	cv::Mat outputs = model.forward();
	std::cout << "output_size = " << outputs.size << std::endl;
	
	// 去除batch维度
	int size[3]{ outputs.size[1], outputs.size[2], outputs.size[3] };
	cv::Mat outputs_1(3, size, CV_32F,outputs.data);
	std::cout << "output1_size = " << outputs_1.size << std::endl;

#pragma region 将每个关键点的热力图提取出来
	for (int i = 0; i < numPoints; ++i)
	{
		cv::Mat singleMat = outputs_1.rowRange(i, i + 1).clone();
		std::cout << singleMat.size << std::endl;

		cv::Mat newMat = singleMat.reshape(1, outputs.size[2]);
		std::cout << newMat.size << std::endl;

		cv::Mat dstMat;
		newMat.convertTo(dstMat, CV_8U, 255.0);
#if 1
		cv::imshow("w", dstMat);
		cv::waitKey();
#endif
		cv::imwrite("./pt_" + std::to_string(i) + ".jpg", dstMat);

		// 找出关键点的位置
		double maxVal, minVal;
		cv::Point maxLoc, minLoc;
		cv::minMaxLoc(dstMat, &minVal, &maxVal, &minLoc, &maxLoc);

#if 1
		cv::circle(resizeImg, maxLoc, 4, colors[i], -1);
#endif

#pragma region 将点还原到原图
		cv::Point outPt;
		keepRatioResizeBackV2(orgSize, inputSize, maxLoc, outPt);
#if 1
		cv::circle(img, outPt, 4, colors[i], -1);
#endif
#pragma endregion

	}
	cv::imshow("keypts in resize img", resizeImg);
	cv::imshow("keypts in org img", img);
	cv::waitKey();
#pragma endregion

	return 0;
}

在这里插入图片描述

吴天德少侠
关注
专栏目录
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

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

抵扣说明：

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

余额充值