VTrain/engine/Step1InferenceEngine/TrainStep1InferenceEngine.cpp

451 lines
17 KiB
C++
Raw Permalink Normal View History

2024-11-27 12:47:45 +00:00
#include "TrainStep1InferenceEngine.h"
#include <opencv2/opencv.hpp>
//#include "myqueue.h"
using namespace ai_matrix;
TrainStep1InferenceEngine::TrainStep1InferenceEngine() {}
TrainStep1InferenceEngine::~TrainStep1InferenceEngine() {}
APP_ERROR TrainStep1InferenceEngine::Init()
{
strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
this->modelConfig_ = Config::getins()->getModelByTrainStep1Config();
this->dataSourceConfig_ = Config::getins()->getDataSourceConfig();
this->identifyConfig_ = Config::getins()->getIdentifyConfig();
int iFolderExist = access(modelConfig_.strModelPath.c_str(), R_OK);
if (iFolderExist == -1)
{
LogError << "模型:" << modelConfig_.strModelPath << " 不存在!";
return false;
}
class_num = this->modelConfig_.vecClass.size();
score_threshold = this->modelConfig_.fScoreThreshold;
int ret = initModel();
if (ret != APP_ERR_OK)
{
LogError << "Failed to read model info, ret = " << ret;
return ret;
}
LogInfo << "Step1InferenceEngine Init ok";
return APP_ERR_OK;
}
APP_ERROR TrainStep1InferenceEngine::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,
this->modelConfig_.strModelPath);
if (nRet != 0)
{
LogError << "YoloV5ClassifyInferenceInit nRet:" << nRet;
return APP_ERR_COMM_READ_FAIL;
}
return APP_ERR_OK;
}
APP_ERROR TrainStep1InferenceEngine::DeInit()
{
yolov5model.YoloV5ClearityInferenceDeinit();
LogInfo << "Step1InferenceEngine DeInit ok";
return APP_ERR_OK;
}
/**
* 1
* inParam : std::vector<stDetection> &vecResult
* outParam: std::vector<stDetection> &vecResult
* return : N/A
*/
void TrainStep1InferenceEngine::getMaxScoreResult(std::vector<stDetection> &vecResult)
{
if (vecResult.size() < 2)
{
return;
}
std::map<Target, std::vector<stDetection>> mapResult;
for (size_t i = 0; i < vecResult.size(); i++)
{
stDetection stDTemp = vecResult.at(i);
if (stDTemp.class_id == 0)
{
mapResult[HEAD].emplace_back(stDTemp);
}
else if (stDTemp.class_id == 1)
{
mapResult[PRO].emplace_back(stDTemp);
}
else if ((stDTemp.class_id >= 2 && stDTemp.class_id <= 6) || stDTemp.class_id == 8 || stDTemp.class_id == 15)
{
mapResult[NUM].emplace_back(stDTemp);
}
else if (stDTemp.class_id >= 9 && stDTemp.class_id <= 17 && stDTemp.class_id != 15)
{
mapResult[TRAINSPACE].emplace_back(stDTemp);
}
else if (stDTemp.class_id == 18)
{
mapResult[SPACE].emplace_back(stDTemp);
}
else if (stDTemp.class_id == 7)
{
mapResult[CONTAINER].emplace_back(stDTemp);
}
}
//清空之前的结果
vecResult.clear();
// 每个类别中,获取得分最高的框
for (auto iter = mapResult.begin(); iter != mapResult.end(); iter++)
{
int iMaxPos = -1;
for (size_t i = 0; i < iter->second.size(); i++)
{
if (iMaxPos == -1)
{
iMaxPos = i;
}
else if (iter->second.at(i).class_conf > iter->second.at(iMaxPos).class_conf)
{
iMaxPos = i;
}
}
if (iMaxPos >= 0)
{
vecResult.emplace_back(iter->second.at(iMaxPos));
}
}
}
/**
*
* inParam : PostSubData &postSubData :
* outParam: PostSubData &postSubData :
* return : N/A
*/
void TrainStep1InferenceEngine::getTargetType(SingleData &singleData)
{
if (singleData.iClassId == TRAIN_HEAD)
{
singleData.iTargetType = HEAD;
}
else if (singleData.iClassId == TRAIN_PRO)
{
singleData.iTargetType = PRO;
}
else if ((singleData.iClassId >= 2 && singleData.iClassId <= 6) ||
singleData.iClassId == J_TRAIN_NUM ||
singleData.iClassId == W_TRAIN_NUM)
{
singleData.iTargetType = NUM;
}
else if (singleData.iClassId >= 9 && singleData.iClassId <= 17 && singleData.iClassId != 15)
{
singleData.iTargetType = TRAINSPACE;
}
else if (singleData.iClassId == U_TRAIN_SPACE)
{
singleData.iTargetType = SPACE;
}
else if (singleData.iClassId == CONTAINERNUM)
{
singleData.iTargetType = CONTAINER;
}
}
/**
*
* inParam : std::vector<stDetection> &vecRet :
: std::shared_ptr<ProcessData> pProcessData :
* outParam: N/A
* return : N/A
*/
void TrainStep1InferenceEngine::filterInvalidInfo(std::vector<stDetection> &vecInferenceResult,
std::shared_ptr<VTrainStep1Data> &pVTrainStep1Data)
{
std::vector<stDetection> vecSpaceInfo;
for (auto it = vecInferenceResult.begin(); it != vecInferenceResult.end();)
{
// LogDebug << " 帧:" << pVTrainStep1Data->iFrameId
// << " --iClassId:" << it->class_id
//// << " iLine:" << it->clear_conf
// << " confidence=" << it->class_conf
// << " lx=" << it->bbox[0]
// << " ly=" << it->bbox[1]
// << " rx=" << it->bbox[2]
// << " ry=" << it->bbox[3]
// << " clear:" << it->clear_conf;
// 根据配置文件中 设置的识别范围,过滤掉无效数据
if (!(it->bbox[0] >= this->dataSourceConfig_.vecIdentifyAreas[0] &&
it->bbox[1] >= this->dataSourceConfig_.vecIdentifyAreas[1] &&
it->bbox[2] <= this->dataSourceConfig_.vecIdentifyAreas[2] &&
it->bbox[3] <= this->dataSourceConfig_.vecIdentifyAreas[3]))
{
LogDebug << "frameId:" << pVTrainStep1Data->iFrameId
<< " 类别:" << it->class_id << " 超出识别区域-识别区域:("
<< this->dataSourceConfig_.vecIdentifyAreas[0] << ","
<< this->dataSourceConfig_.vecIdentifyAreas[1] << "),("
<< this->dataSourceConfig_.vecIdentifyAreas[2] << ","
<< this->dataSourceConfig_.vecIdentifyAreas[2] << ")";
it = vecInferenceResult.erase(it);
continue;
}
// 如果设置了不识别车头,则去掉车头标记的大框 !this->identifyConfig_.bTrainHeardDetect &&
if (it->class_id == TRAIN_HEAD)
{
LogDebug << "frameId:" << pVTrainStep1Data->iFrameId << " 过滤掉车头编号";
it = vecInferenceResult.erase(it);
continue;
}
// 去除车头时的非车头编号信息
if(pVTrainStep1Data->iTrainStage == MONITOR_MODEL_TRAIN_HEAD )
{
if(it->class_id != TRAIN_HEAD)
{
LogDebug << " 帧号:" << pVTrainStep1Data->iFrameId
<< " 大类:" << it->class_id << " 识别于车头位置,无效!";
it = vecInferenceResult.erase(it);
continue;
}
}
// 去除车尾的车头编号信息
if (pVTrainStep1Data->iTrainStage != MONITOR_MODEL_TRAIN_HEAD)
{
if (it->class_id == TRAIN_HEAD)
{
LogDebug << " 帧号:" << pVTrainStep1Data->iFrameId
<< " 大类:" << it->class_id << " 识别于非车头位置,无效!";
it = vecInferenceResult.erase(it);
continue;
}
}
// 去除车尾 和 车头车体之间 的间隔信息
if ((pVTrainStep1Data->iTrainStage == MONITOR_MODEL_TRAIN_TAIL || pVTrainStep1Data->iTrainStage == MONITOR_MODEL_HEAD_FIRST)
&& (it->class_id >= C_TRAIN_SPACE && it->class_id <= U_TRAIN_SPACE && it->class_id != W_TRAIN_NUM))
{
LogDebug << " frameId:" << pVTrainStep1Data->iFrameId
<< " bigclassid:" << it->class_id
<<" 识别于车尾或者车头与车身交接部分,无效!";
it = vecInferenceResult.erase(it);
continue;
}
// 过滤掉识别于模型反馈无车状态下的所有大框信息
if (pVTrainStep1Data->iTrainStage == MONITOR_MODEL_NO_TRAIN)
{
LogDebug << " frameId:" << pVTrainStep1Data->iFrameId
<< " bigclassid:" << it->class_id
<<" 识别于模型反馈的无车状态下,无效!";
it = vecInferenceResult.erase(it);
continue;
}
//剔除高度大于宽的车号大框
if (((it->class_id >= K_TRAIN_NUM && it->class_id <= NX_TRAIN_NUM)
|| it->class_id == J_TRAIN_NUM
|| it->class_id == W_TRAIN_NUM) &&
(it->bbox[3] - it->bbox[1]) > (it->bbox[2] - it->bbox[0]))
{
LogWarn << " frameId:" << pVTrainStep1Data->iFrameId
<< " bigclassid:" << it->class_id << " 过滤 高度大于宽度的车号";
it = vecInferenceResult.erase(it);
continue;
}
if (it->class_id == K_TRAIN_NUM)
{
int iCenterY = IMAGE_HEIGHT / 2;
int iHeight0 = it->bbox[1] / 2 + it->bbox[3] / 2;
if (iHeight0 > iCenterY) {
LogWarn << "矿车编号大框在画面Y轴中线以下帧号:"
<< pVTrainStep1Data->iFrameId
<< " 画面Y轴中心" << iCenterY
<< " 大框Y轴中心" << iHeight0 ;
// << "[" << it->bbox[0] << "," << it->bbox[1] << "]"
// << "[" << it->bbox[2] << "," << it->bbox[3] << "]";
it = vecInferenceResult.erase(it);
continue;
}
}
if (it->class_id >= C_TRAIN_SPACE && it->class_id <= U_TRAIN_SPACE && it->class_id != W_TRAIN_NUM)
{
vecSpaceInfo.emplace_back(*it);
}
++it;
}
if (vecInferenceResult.size() <= 0) return;
// 过滤与间隔X轴重合的其他大框
// for (auto it = vecInferenceResult.begin(); it != vecInferenceResult.end();)
// {
// if (!((it->class_id >= 9 && it->class_id <= 17 && it->class_id != 15) || it->class_id == 18))
// {
// for (int i = 0; i < vecSpaceInfo.size(); i++)
// {
// if (
// (it->bbox[0] > vecSpaceInfo[i].bbox[0]
// && it->bbox[0] < vecSpaceInfo[i].bbox[2])
// ||
// (it->bbox[2] > vecSpaceInfo[i].bbox[0]
// && it->bbox[2] < vecSpaceInfo[i].bbox[2])
// )
// {
// LogWarn << "-- " << it->class_id << " _ " << vecSpaceInfo[i].class_id
// << " " << it->bbox[0] << "," << it->bbox[2] << " || " << vecSpaceInfo[i].bbox[0] << "," << vecSpaceInfo[i].bbox[2];
// it = vecInferenceResult.erase(it);
// break;
// }
// }
// }
// }
//主摄像头1帧如果只识别2个大框如果非平车的车号和属性场景则必有间隔框
if (vecInferenceResult.size() >= 2)
{
int iHeight0 = vecInferenceResult[0].bbox[1] / 2 + vecInferenceResult[0].bbox[3] / 2;
int iHeight1 = vecInferenceResult[1].bbox[1] / 2 + vecInferenceResult[1].bbox[3] / 2;
int iCenterY = IMAGE_HEIGHT / 2;
if (iHeight0 < iCenterY && iHeight1 < iCenterY) //非平车
{
bool bHaveSpace = false;
for (auto &it : vecInferenceResult)
{
if (it.class_id >= C_TRAIN_SPACE
&& it.class_id <= U_TRAIN_SPACE
&& it.class_id != W_TRAIN_NUM)
{
bHaveSpace = true;
}
}
if (!bHaveSpace)
{
LogDebug << " frameId:" << pVTrainStep1Data->iFrameId << " no space";
vecInferenceResult.clear();
}
// if (!(vecInferenceResult[0].class_id >= C_TRAIN_SPACE
// && vecInferenceResult[0].class_id <= U_TRAIN_SPACE
// && vecInferenceResult[0].class_id != W_TRAIN_NUM)
// && !(vecInferenceResult[1].class_id >= C_TRAIN_SPACE
// && vecInferenceResult[1].class_id <= U_TRAIN_SPACE
// && vecInferenceResult[1].class_id != W_TRAIN_NUM))
// {
// LogDebug << " frameId:" << pVTrainStep1Data->iFrameId << " no space";
// vecInferenceResult.clear();
// }
}
}
}
APP_ERROR TrainStep1InferenceEngine::Process()
{
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<VTrainStep1Data> pVTrainStep1Data = std::static_pointer_cast<VTrainStep1Data>(pVoidData0);
if (pVTrainStep1Data->cvImage.empty())
{
usleep(1000); //1ms
continue;
}
// else
// {
// vector<int> compression_params;
// compression_params.push_back(cv::IMWRITE_JPEG_QUALITY); //选择jpeg
// compression_params.push_back(100); //图片质量
// cv::imwrite("./jpg/" + std::to_string(pVTrainStep1Data->iFrameId) + ".jpg", pVTrainStep1Data->cvImage, compression_params);
//
// }
//进行推理
std::vector<stDetection> vecInferenceResult;
yolov5model.YoloV5ClearityInferenceModel(pVTrainStep1Data->cvImage, vecInferenceResult);
//过滤无效信息
this->filterInvalidInfo(vecInferenceResult, pVTrainStep1Data);
this->getMaxScoreResult(vecInferenceResult);
std::shared_ptr<InferenceResultData> pInferenceResultData = std::make_shared<InferenceResultData>();
pInferenceResultData->iFrameId = pVTrainStep1Data->iFrameId;
pInferenceResultData->bIsEnd = pVTrainStep1Data->bIsEnd;
pInferenceResultData->strTrainDate = pVTrainStep1Data->strTrainDate;
pInferenceResultData->strTrainTime = pVTrainStep1Data->strTrainTime;
for (size_t j = 0; j < vecInferenceResult.size(); j++)
{
/*
[0:; 1:; 2:(,); 3:; 4:; 5:; 6:
7:; 8:; 9:; 10:; 11:; 12:; 13:;
14:; 15:; 16: ; 17:; 18:; 19:; 20:]
*/
if (vecInferenceResult[j].class_id < 0 || vecInferenceResult[j].class_id > 20)
{
continue;
}
SingleData singledata;
singledata.iClassId = vecInferenceResult[j].class_id;
singledata.fScore = vecInferenceResult[j].class_conf;
singledata.fLTX = vecInferenceResult[j].bbox[0];
singledata.fLTY = vecInferenceResult[j].bbox[1];
singledata.fRBX = vecInferenceResult[j].bbox[2];
singledata.fRBY = vecInferenceResult[j].bbox[3];
singledata.fClear = vecInferenceResult[j].clear_id;
this->getTargetType(singledata);
pInferenceResultData->vecSingleData.emplace_back(singledata);
LogDebug << " 帧:" << pInferenceResultData->iFrameId
<< " --iClassId:" << singledata.iClassId
<< " iLine:" << singledata.iLine
<< " confidence=" << singledata.fScore
<< " lx=" << singledata.fLTX
<< " ly=" << singledata.fLTY
<< " rx=" << singledata.fRBX
<< " ry=" << singledata.fRBY
<< " clear:" << singledata.fClear;
}
outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pInferenceResultData), true);
}
return APP_ERR_OK;
}