Wafer_Inspection/Code/MF_TRTInfer/MA_TRTInferAlgoBase.cpp

#include "MA_TRTInferAlgoBase.h"
#include "ML_Log.h"
#include <time.h>

MA_TRTInferAlgoBase::MA_TRTInferAlgoBase(const trtUtils::InitParameter& param)
{
	mStream = nullptr;
	m_param = param;

	// ³õʼ»¯ÈÕÖ¾
	mLogPtr = std::make_shared<ML_Log>();
	std::string logFile = "./logFiles/vision_log-20240418.txt";

	if (mLogPtr)
	{
		mLogPtr->Init(logFile.c_str());
	}

}

MA_TRTInferAlgoBase::~MA_TRTInferAlgoBase()
{
	mLogPtr->UnInit();
}

bool MA_TRTInferAlgoBase::initEngine(const std::string& _onnxFileName)
{
	return true;
}

bool MA_TRTInferAlgoBase::check()
{
	if (m_engine == nullptr || m_context == nullptr)
	{
		LOG_ERROR("do check, engine or context dose not engine. \n");
		return false;
	}

	int idx;
	nvinfer1::Dims dims;

	//sample::gLogInfo << "the engine's info:" << std::endl;
	LOG_INFO("The engine's info:");
	for (auto layer_name : m_param.input_output_names)
	{
		idx = m_engine->getBindingIndex(layer_name.c_str());
		dims = m_context->getBindingDimensions(idx);
		sample::gLogInfo << "idx = " << idx << ", " << layer_name << ": ";
		LOG_INFO("idx = " + std::to_string(idx) + ", " + layer_name);
		for (int i = 0; i < dims.nbDims; i++)
		{
			sample::gLogInfo << dims.d[i] << ", ";
		}
		sample::gLogInfo << std::endl;
		LOG_INFO("\n");
	}
	sample::gLogInfo << "the context's info:" << std::endl;
	LOG_INFO("The context's info:");
	for (auto layer_name : m_param.input_output_names)
	{
		idx = m_engine->getBindingIndex(layer_name.c_str());
		dims = m_context->getBindingDimensions(idx);
		sample::gLogInfo << "idx = " << idx << ", " << layer_name << ": ";
		for (int i = 0; i < dims.nbDims; i++)
		{
			sample::gLogInfo << dims.d[i] << ", ";
		}
		sample::gLogInfo << std::endl;
	}

	return true;
}

bool MA_TRTInferAlgoBase::doTRTInfer(const std::vector<MN_VisionImage::MS_ImageParam>& _bufImgs, std::vector<trtUtils::MR_Result>* _detectRes, int* _user)
{
	return true;
}

bool MA_TRTInferAlgoBase::doTRTInfer(const std::vector<cv::Mat>& _bufImgs, std::vector<trtUtils::MR_Result>* _detectRes, int* _user)
{
	return true;
}

std::string MA_TRTInferAlgoBase::getError()
{
	return " ";
}

void MA_TRTInferAlgoBase::freeMemeory()
{
	return;
}

bool MA_TRTInferAlgoBase::measureAxis(std::vector<double>& measureRes, const MN_VisionImage::MS_ImageParam& _bufImg)
{
	return true;
}

int MA_TRTInferAlgoBase::loadTRTModelData(const std::string& _trtFile, std::vector<uchar>& _modelData)
{
	std::ifstream in(_trtFile, std::ios::in | std::ios::binary);
	if (!in.is_open())
	{
		return {};
	}

	in.seekg(0, std::ios::end);
	size_t len = in.tellg();

	if (len > 0)
	{
		in.seekg(0, std::ios::beg);
		_modelData.resize(len);

		in.read((char*)&_modelData[0], len);
	}
	in.close();

	return 0;
}

int MA_TRTInferAlgoBase::copyToDevice(const std::vector<cv::Mat>& _imgsBatch)
{
	return 0;
}

int MA_TRTInferAlgoBase::preProcess(const std::vector<cv::Mat>& _imgsBatch)
{
	return 0;
}

int MA_TRTInferAlgoBase::infer()
{
	return 0;
}

int MA_TRTInferAlgoBase::copyFromDevice(const std::vector<cv::Mat>& _imgsBatch)
{
	return 0;
}

int MA_TRTInferAlgoBase::postProcess(const std::vector<cv::Mat>& _imgsBatch)
{
	return 0;
}

bool MA_TRTInferAlgoBase::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;
}