Python —— Windows10下通过Opencv-dnn推理Yolov5的分割模型（pt转onnx）（附源码）

附：Python —— Windows10下训练Yolov5分割模型并测试

 

视频效果

     手机拍摄一段工位视频，上传到win10训练了yolov5分割鼠标的样本后，将训练出的pt模型转到onnx模型，通过opencv_dnn模块进行实时推理效果。

     下面视频演示了python用torch推理效果、opencv-dnn模块推理效果:

 

pt转onnx

像yolov5目标检测时将pt转onnx一样，修改调用源码根目录export.py即可。

 

源码

#include 

#include 
#include 
#include 

std::vector<std::string> class_list;

// INPUT_WIDTH、INPUT_HEIGHT必须与训练时设置的尺寸一致
const float INPUT_WIDTH = 640.0;
const float INPUT_HEIGHT = 640.0;

struct OutputSeg 
{
   
	int id;             //结果类别id
	float confidence;   //结果置信度
	cv::Rect box;       //矩形框
	cv::Mat boxMask;       //矩形框内mask，节省内存空间和加快速度
};

static struct MaskParams 
{
   
	int segChannels = 32;
	int segWidth = 160;
	int segHeight = 160;
	int netWidth = 640;
	int netHeight = 640;
	float maskThreshold = 0.5;
	cv::Size srcImgShape;
	cv::Vec4d params;

};

void LetterBox(const cv::Mat& image, cv::Mat& outImage, cv::Vec4d& params, 
	const cv::Size& newShape= cv::Size(INPUT_WIDTH, INPUT_HEIGHT),
	bool autoShape = false, 
	bool scaleFill = false, 
	bool scaleUp = true, 
	int stride = 32, 
	const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
   
	if (false)
	{
   
		int maxLen = MAX(image.rows, image.cols);
		outImage = cv::Mat::zeros(cv::Size(maxLen, maxLen), CV_8UC3);
		image.copyTo(outImage(cv::Rect(0, 0, image.cols, image.rows)));
		params[0] = 1;
		params[1] = 1;
		params[3] = 0;
		params[2] = 0;
	}

	cv::Size shape = image.size();
	float r = std::min((float)newShape.height / (float)shape.height,
		(float)newShape.width / (float)shape.width);
	if (!scaleUp)
		r = std::min(r, 1.0f);

	float ratio[2]{
    r, r };
	int new_un_pad[2] = {
    (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };

	auto dw = (float)(newShape.width - new_un_pad[0]);
	auto dh = (float)(newShape.height - new_un_pad[1]);

	if (autoShape)
	{
   
		dw = (float)((int)dw % stride);
		dh = (float)((int)dh %