Train_Identify/nvidia_ascend_engine/nvidia_engine/TrainStepTwoEngine/TrainStepTwoEngine.cpp

247 lines
9.2 KiB
C++
Raw Normal View History

2024-01-23 02:46:26 +00:00
#include "TrainStepTwoEngine.h"
#include <opencv2/opencv.hpp>
#include "myutils.h"
#include "myqueue.h"
using namespace ai_matrix;
TrainStepTwoEngine::TrainStepTwoEngine() {}
TrainStepTwoEngine::~TrainStepTwoEngine() {}
APP_ERROR TrainStepTwoEngine::Init()
{
bUseEngine_ = MyUtils::getins()->ChkIsHaveTarget("NUM");
if (!bUseEngine_)
{
LogWarn << "engineId_:" << engineId_ << " not use engine";
return APP_ERR_OK;
}
strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
modelConfig_ = MyYaml::GetIns()->GetModelConfig("TrainStepTwoEngine");
//读取模型信息
APP_ERROR ret = ReadModelInfo();
if (ret != APP_ERR_OK)
{
LogError << "Failed to read model info, ret = " << ret;
return ret;
}
ret = InitModel();
if (ret != APP_ERR_OK)
{
LogError << "Failed to read model info, ret = " << ret;
return ret;
}
LogInfo << "AclTrainStepTwoEngine Init ok";
return APP_ERR_OK;
}
APP_ERROR TrainStepTwoEngine::InitModel()
{
modelinfo.yolov5ClearityModelParam.uiClassNum = class_num;
modelinfo.yolov5ClearityModelParam.uiClearNum = clear_num;
modelinfo.yolov5ClearityModelParam.uiDetSize = det_size;
modelinfo.yolov5ClearityModelParam.fScoreThreshold = score_threshold;
modelinfo.yolov5ClearityModelParam.fNmsThreshold = nms_threshold;
modelinfo.modelCommonInfo.uiModelWidth = model_width;
modelinfo.modelCommonInfo.uiModelHeight = model_height;
modelinfo.modelCommonInfo.uiInputSize = input_size;
modelinfo.modelCommonInfo.uiOutputSize = output_size;
modelinfo.modelCommonInfo.uiChannel = INPUT_CHANNEL;
modelinfo.modelCommonInfo.uiBatchSize = batch_size;
modelinfo.modelCommonInfo.strInputBlobName = INPUT_BLOB_NAME;
modelinfo.modelCommonInfo.strOutputBlobName = OUTPUT_BLOB_NAME;
string strModelName = "";
int nRet = yolov5model.YoloV5ClearityInferenceInit(&modelinfo, strModelName, modelConfig_.strOmPath);
if (nRet != 0)
{
LogInfo << "YoloV5ClassifyInferenceInit nRet:" << nRet;
return APP_ERR_COMM_READ_FAIL;
}
return APP_ERR_OK;
}
APP_ERROR TrainStepTwoEngine::ReadModelInfo()
{
char szAbsPath[PATH_MAX];
// Get the absolute path of model file
if (realpath(modelConfig_.strOmPath.c_str(), szAbsPath) == nullptr)
{
LogError << "Failed to get the real path of " << modelConfig_.strOmPath.c_str();
return APP_ERR_COMM_NO_EXIST;
}
// Check the validity of model path
int iFolderExist = access(szAbsPath, R_OK);
if (iFolderExist == -1)
{
LogError << "ModelPath " << szAbsPath << " doesn't exist or read failed!";
return APP_ERR_COMM_NO_EXIST;
}
//读取模型参数信息文件
Json::Value jvModelInfo;
if (!MyUtils::getins()->ReadJsonInfo(jvModelInfo, modelConfig_.strModelInfoPath))
{
LogError << "ModelInfoPath:" << modelConfig_.strModelInfoPath << " doesn't exist or read failed!";
return APP_ERR_COMM_NO_EXIST;
}
model_width = jvModelInfo["model_width"].asInt();
model_height = jvModelInfo["model_height"].asInt();
clear_num = jvModelInfo["clear"].isArray() ? jvModelInfo["clear"].size() : 0;
class_num = jvModelInfo["class"].isArray() ? jvModelInfo["class"].size() : 0;
input_size = GET_INPUT_SIZE(model_width, model_height);
output_size = GET_OUTPUT_SIZE(model_width, model_height, clear_num, class_num);
det_size = clear_num + class_num + 5;
score_threshold = modelConfig_.fScoreThreshold;
nms_threshold = modelConfig_.fNMSTreshold;
return APP_ERR_OK;
}
APP_ERROR TrainStepTwoEngine::DeInit()
{
if (!bUseEngine_)
{
LogWarn << "engineId_:" << engineId_ << " not use engine";
return APP_ERR_OK;
}
yolov5model.YoloV5ClearityInferenceDeinit();
LogInfo << "TrainStepTwoEngine DeInit ok";
return APP_ERR_OK;
}
/**
* push数据到队列push
* inParam : const std::string strPort push的端口
: const std::shared_ptr<ProcessData> &pProcessData push的数据
* outParam: N/A
* return : N/A
*/
void TrainStepTwoEngine::PushData(const std::string &strPort, const std::shared_ptr<ProcessData> &pProcessData)
{
while (true)
{
int iRet = outputQueMap_[strPort]->push(std::static_pointer_cast<void>(pProcessData));
if (iRet != 0)
{
LogDebug << "sourceid:" << pProcessData->iDataSource << " frameid:" << pProcessData->iFrameId << " push fail iRet:" << iRet;
if (iRet == 2)
{
usleep(10000); // 10ms
continue;
}
}
break;
}
}
APP_ERROR TrainStepTwoEngine::Process()
{
if (!bUseEngine_)
{
LogWarn << "engineId_:" << engineId_ << " not use engine";
return APP_ERR_OK;
}
int iRet = APP_ERR_OK;
while (!isStop_)
{
std::shared_ptr<void> pVoidData0 = nullptr;
inputQueMap_[strPort0_]->pop(pVoidData0);
if (nullptr == pVoidData0)
{
usleep(1000); //1ms
continue;
}
std::shared_ptr<ProcessData> pProcessData = std::static_pointer_cast<ProcessData>(pVoidData0);
//组织输出数据
std::shared_ptr<PostData> pPostData = std::make_shared<PostData>();
pPostData->iModelType = MODELTYPE_NUM;
//获取图片
if (pProcessData->iStatus == TRAINSTATUS_RUN || pProcessData->bIsEnd)
{
if (pProcessData->pData != nullptr && pProcessData->iSize != 0)
{
std::shared_ptr<PostData> ppostbuff = std::static_pointer_cast<PostData>(pProcessData->pVoidData);
cv::Mat img(pProcessData->iHeight, pProcessData->iWidth, CV_8UC3, static_cast<uint8_t *>(pProcessData->pData.get())); //RGB
for(int i = 0; i< ppostbuff->vecPostSubData.size(); i++)
{
PostSubData postsubdata = ppostbuff->vecPostSubData[i];
if (postsubdata.iTargetType != NUM && postsubdata.iTargetType != PRO && postsubdata.iTargetType != HEAD)
{
continue;
}
cv::Rect step2_rect(cv::Point(postsubdata.step1Location.fLTX, postsubdata.step1Location.fLTY), cv::Point(postsubdata.step1Location.fRBX, postsubdata.step1Location.fRBY));
cv::Mat step2_image = img(step2_rect).clone();
//进行推理
std::vector<stDetection> res;
auto start = std::chrono::system_clock::now(); // 计时开始
yolov5model.YoloV5ClearityInferenceModel(step2_image, res);
auto end = std::chrono::system_clock::now();
LogInfo << "nopr2 inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms";
PostSubData postSubDataNew;
postSubDataNew.iTargetType = postsubdata.iTargetType;
postSubDataNew.iBigClassId = postsubdata.iBigClassId;
postSubDataNew.iCarXH = postsubdata.iCarXH;
postSubDataNew.step1Location = postsubdata.step1Location;
//整理推理结果
//根据非极大值抑制的结果标注相关信息(画框,文字信息等)
//res.size()为每张图片上的识别到的对象数目
for (size_t j = 0; j < res.size(); j++)
{
SingleData singledata;
singledata.iLine = res[j].clear_id;
// singledata.iLine = -1;
singledata.iClassId = res[j].class_id;
singledata.fScore = res[j].class_conf;
// singledata.iAnchorId = -1;
singledata.fLTX = res[j].bbox[0];
singledata.fLTY = res[j].bbox[1];
singledata.fRBX = res[j].bbox[2];
singledata.fRBY = res[j].bbox[3];
singledata.fClear = res[j].clear_id;
MyUtils::getins()->Step2ResetLocation(singledata, 1.0, pProcessData, postSubDataNew.step1Location);
postSubDataNew.vecSingleData.emplace_back(singledata);
LogDebug << "sourceid:" << pProcessData->iDataSource << " step2 after frameId:" << pProcessData->iFrameId
<< " --iClassId:" << singledata.iClassId << " iLine:" << singledata.iLine << " confidence=" << singledata.fScore
<< " lx=" << singledata.fLTX << " ly=" << singledata.fLTY << " rx=" << singledata.fRBX << " ry=" << singledata.fRBY;
}
pPostData->vecPostSubData.emplace_back(postSubDataNew);
}
}
}
//及时释放内存
if (pProcessData->pData != nullptr)
{
pProcessData->pData = nullptr;
pProcessData->iSize = 0;
}
// push端口0第1步推理
pProcessData->pVoidData = std::static_pointer_cast<void>(pPostData);
PushData(strPort0_, pProcessData);
}
return APP_ERR_OK;
}