Wafer_Inspection/Code/TestTRTInterDll/TestTRTInterDll.cpp

// TestTRTInterDll.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include <iostream>
#include <filesystem>
#include <iostream>
#include <string>
#include "utils.h"
#include "MI_Interface.h"
#include <future>


namespace fs = std::filesystem;

bool buffer2Mat(const MN_VisionImage::MS_ImageParam& _inImg, cv::Mat& _mat);

void test01()
{
	std::vector<std::string> imagePath;
	imagePath.emplace_back("./imageDatas/daisy.jpg");
	imagePath.emplace_back("./imageDatas/dandelion.jpg");
	imagePath.emplace_back("./imageDatas/sunflower.jpg");
	imagePath.emplace_back("./imageDatas/tulip.jpg");

	std::vector<cv::Mat> matImgs;
	for (auto _var : imagePath)
	{
		matImgs.emplace_back(cv::imread(_var, 1));
	}

	trtUtils::InitParameter params;
	params.m_modelType = trtUtils::ME_ModelType::E_RESNET34;
	params.class_names = trtUtils::dataSets::flower_labels;
	//initParameters.class_names = utils::dataSets::voc20;
	params.num_class = 5; // for flower_labels
	//initParameters.num_class = 20; // for voc2012
	params.batch_size = 1;
	params.dst_h = 224;
	params.dst_w = 224;
	memcpy(&params.mImage.m_data, matImgs[0].data, sizeof(matImgs[0].rows * matImgs[0].cols * matImgs[0].channels()));
	params.mImage.m_height = matImgs[0].rows;
	params.mImage.m_width = matImgs[0].cols;
	params.mImage.m_channels = 3;
	params.input_output_names = { "input",  "output" };
	params.conf_thresh = 0.25f;
	params.iou_thresh = 0.45f;
	params.save_path = "./imageDatas";
	params.meanVec = { 0.406, 0.456, 0.485 };
	params.stdVec = { 0.225, 0.224, 0.229 };
	MI_VisionInterface* resnet34Ptr = getInterfacePtr(params);
	if (!resnet34Ptr)
	{
		return;
	}

	trtUtils::InitParameter params1;
	params1.m_modelType = trtUtils::ME_ModelType::E_RESNET50;
	params1.class_names = trtUtils::dataSets::flower_labels;
	//initParameters.class_names = utils::dataSets::voc20;
	params1.num_class = 5; // for flower_labels
	//initParameters.num_class = 20; // for voc2012
	params1.batch_size = 1;
	params1.dst_h = 224;
	params1.dst_w = 224;
	memcpy(&params1.mImage.m_data, matImgs[0].data, sizeof(matImgs[0].rows * matImgs[0].cols * matImgs[0].channels()));
	params1.mImage.m_height = matImgs[0].rows;
	params1.mImage.m_width = matImgs[0].cols;
	params1.mImage.m_channels = 3;
	params1.input_output_names = { "input",  "output" };
	params1.conf_thresh = 0.25f;
	params1.iou_thresh = 0.45f;
	params1.save_path = "./imageDatas";
	params1.meanVec = { 0.406, 0.456, 0.485 };
	params1.stdVec = { 0.225, 0.224, 0.229 };
	MI_VisionInterface* resnet50Ptr = getInterfacePtr(params1);
	if (!resnet50Ptr)
	{
		return;
	}

	// 线程测试，初始化在一个线程，推理在另外一个线程内
	std::string onnxFile = "./imageDatas/resnet34_0407.onnx";
	std::string onnxFile1 = "./imageDatas/resnet50.onnx";

#if 0
	bool bRet = false;
	bRet = resnetPtr->initEngine(onnxFile);
	if (!bRet)
	{
		return;
	}

	bRet = resnetPtr->check();
	if (!bRet)
	{
		return;
	}

	std::vector<MR_Result> detectResVec;
	bRet = resnetPtr->doTRTInfer(matImgs, &detectResVec, nullptr);
	if (!bRet)
	{
		return;
	}
#endif // 0

	// 初始化线程
	auto initRes = std::async(std::launch::async, [&] {

		bool bRet = false;
		// 加载第一个模型初始化并推理
		bRet = resnet34Ptr->initEngine(onnxFile);
		if (!bRet)
		{
			return false;
		}

		// 加载第二个模型初始化并推理
		bRet = resnet50Ptr->initEngine(onnxFile1);
		if (!bRet)
		{
			return false;
		}

		return true;
		});
	initRes.wait();

	auto initRes1 = std::async(std::launch::async, [&] {

		bool bRet = false;
		bRet = resnet34Ptr->check();
		if (!bRet)
		{
			return false;
		}

		std::vector<trtUtils::MR_Result> detectResVec1;
		bRet = resnet34Ptr->doTRTInfer(matImgs, &detectResVec1, nullptr);
		if (!bRet)
		{
			return false;
		}

		bRet = resnet50Ptr->check();
		if (!bRet)
		{
			return false;
		}

		std::vector<trtUtils::MR_Result> detectResVec2;
		bRet = resnet50Ptr->doTRTInfer(matImgs, &detectResVec2, nullptr);
		if (!bRet)
		{
			return false;
		}

		return true;

		});

	//initRes1.wait();

	auto initRes2 = std::async(std::launch::async, [&] {

		bool bRet = false;
		bRet = resnet34Ptr->check();
		if (!bRet)
		{
			return false;
		}

		std::vector<trtUtils::MR_Result> detectResVec1;
		bRet = resnet34Ptr->doTRTInfer(matImgs, &detectResVec1, nullptr);
		if (!bRet)
		{
			return false;
		}

		bRet = resnet50Ptr->check();
		if (!bRet)
		{
			return false;
		}

		std::vector<trtUtils::MR_Result> detectResVec2;
		bRet = resnet50Ptr->doTRTInfer(matImgs, &detectResVec2, nullptr);
		if (!bRet)
		{
			return false;
		}

		return true;
		});

	//initRes2.wait();

}

void test02()
{
	cv::Mat img = cv::imread("./imageDatas/daisy.jpg", 1);
	MN_VisionImage::MS_ImageParam bufImg((uchar*)img.data, img.cols, img.rows, MN_VisionImage::ME_ImageType::E_RGB);

	cv::Mat convertImg;
	bool bRet = buffer2Mat(bufImg, convertImg);
	cv::imshow("src", img);
	cv::imshow("image", convertImg);
	cv::waitKey(0);
}

int main()
{

	test01();

	system("pause");
	return 0;
}

// 运行程序: Ctrl + F5 或调试 >“开始执行(不调试)”菜单
// 调试程序: F5 或调试 >“开始调试”菜单

bool buffer2Mat(const MN_VisionImage::MS_ImageParam& _inImg, cv::Mat& _mat)
{
	uchar* pBuf = _inImg.m_data.get(); //获取图像数据首地址
	int nW = _inImg.m_width;
	int nH = _inImg.m_height;
	int nChannel = _inImg.m_channels;

	if (pBuf == nullptr || nW <= 1 || nH <= 1)
	{
		// LOG_ERROR("convert buffer to mat, input image error. \n");
		return false;
	}

	if (_inImg.mImgType == MN_VisionImage::ME_ImageType::E_GRAY)
	{
		_mat = cv::Mat(nH, nW, CV_8UC1, pBuf);
	}
	else if (_inImg.mImgType == MN_VisionImage::ME_ImageType::E_RGBA)
	{
		_mat = cv::Mat(nH, nW, CV_8UC4, pBuf);
	}
	else
	{
		_mat = cv::Mat(nH, nW, CV_8UC3, pBuf);
	}

	if (_mat.data == nullptr || _mat.cols <= 1 || _mat.rows <= 1)
	{
		// LOG_ERROR("convert buffer to mat, convert image failed. \n");
		return false;
	}

	return true;
}


// 0.3 * 0.3 - 12寸晶圆， 一次性识别9颗 -- 测试推理时间
/*
1. 芯片分类：正常，边缘，蓝膜
2. 对于正常的芯片进行检测--针对划痕和脏污分两个模型/或者一个模型；
3. 针对针测点检测 -- 传统算法；
*/