generated from zhangwei/Train_Identify
123 lines
5.4 KiB
C++
123 lines
5.4 KiB
C++
#include "PostProcessEngine.h"
|
|
|
|
using namespace std;
|
|
using namespace ai_matrix;
|
|
|
|
const std::vector<std::string> class_names={
|
|
"person", "bicycle", "car", "motorcycle", "airplane",
|
|
"bus", "train", "truck", "boat", "traffic light", "fire hydrant",
|
|
"stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
|
|
"sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
|
|
"umbrella", "handbag", "tie", "suitcase", "frisbee", "skis",
|
|
"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
|
|
"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
|
|
"cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich",
|
|
"orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
|
|
"chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv",
|
|
"laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
|
|
"oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
|
|
"scissors", "teddy bear", "hair drier", "toothbrush"
|
|
};
|
|
|
|
|
|
PostProcessEngine::PostProcessEngine() {}
|
|
PostProcessEngine::~PostProcessEngine() {}
|
|
|
|
APP_ERROR PostProcessEngine::Init()
|
|
{
|
|
strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
|
|
|
|
width_ = IMAGE_WIDTH, height_ = IMAGE_HEIGHT;
|
|
|
|
LogInfo << "engineId_:" << engineId_ << " ImagePreprocessEngine Init ok";
|
|
return APP_ERR_OK;
|
|
}
|
|
|
|
APP_ERROR PostProcessEngine::DeInit()
|
|
{
|
|
LogInfo << "engineId_:" << engineId_ << " ImagePreprocessEngine DeInit ok";
|
|
return APP_ERR_OK;
|
|
}
|
|
|
|
|
|
APP_ERROR PostProcessEngine::Process()
|
|
{
|
|
uint64_t u64image_num = 0;
|
|
std::string result;
|
|
|
|
uint64_t u64count_num = 0;
|
|
int iRet = APP_ERR_OK;
|
|
|
|
while (!isStop_)
|
|
{
|
|
//后处理
|
|
std::shared_ptr<void> pVoidData0 = nullptr;
|
|
inputQueMap_[strPort0_]->pop(pVoidData0);
|
|
if (nullptr == pVoidData0)
|
|
{
|
|
usleep(1*1000); //n ms
|
|
continue;
|
|
}
|
|
|
|
// std::cout<<"Enter PostProcessEngine Thread "<<++u64count_num<<" Times!"<<std::endl;
|
|
std::shared_ptr<InferenceData> pInferenceModelData = std::static_pointer_cast<InferenceData>(pVoidData0);
|
|
|
|
std::vector<Yolo::Detection> res;
|
|
nms(res, static_cast<float *>(pInferenceModelData->pData.get()), CONF_THRESH, NMS_THRESH); //进行非极大值抑制后res中保留着结果
|
|
std::cout<<"this frame find obj num is:"<<res.size()<<std::endl;
|
|
//根据非极大值抑制的结果标注相关信息(画框,文字信息等)
|
|
//res.size()为每张图片上的识别到的对象数目
|
|
cv::Mat RGBImage(height_, width_, CV_8UC3, pInferenceModelData->pSrcData.get()); //RGB源数据
|
|
|
|
// for (size_t j = 0; j < res.size(); j++) {
|
|
// cv::Rect r = get_rect(RGBImage, res[j].bbox);
|
|
// cv::rectangle(RGBImage, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
|
|
// cv::putText(RGBImage, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
|
|
// }
|
|
// result = "postprocess_result_"+to_string(++u64image_num)+".jpg";
|
|
// cv::imwrite(result, RGBImage); //Jpeg库与NVIDIA底层库有冲突,使用硬件Jpeg编码,使用OpenCV也不行,会报段错误!!!
|
|
|
|
for (size_t j = 0; j < res.size(); j++) {
|
|
uint8_t r, g, b;
|
|
tie(r, g, b) = random_color(res[j].class_id);
|
|
std::string put_text = class_names.at(res[j].class_id) + ": " + std::to_string(res[j].conf);
|
|
cv::Rect rect = get_rect(RGBImage, res[j].bbox);
|
|
cv::rectangle(RGBImage, rect, cv::Scalar(r, g, b), 5);
|
|
// cv::putText(RGBImage, put_text, cv::Point(rect.x, rect.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(r, g, b), 2);
|
|
|
|
int width = cv::getTextSize(put_text, 0, 1, 2, nullptr).width + 10;
|
|
cv::rectangle(RGBImage, cv::Point(rect.x-3, rect.y-33), cv::Point(rect.x + width, rect.y), cv::Scalar(r, g, b), -1);
|
|
cv::putText(RGBImage, put_text, cv::Point(rect.x, rect.y-5), 0, 1, cv::Scalar::all(0), 2, 16);
|
|
}
|
|
|
|
//构造画框后的RGB数据
|
|
void* pPostProcessRGBBuffer = nullptr;
|
|
unsigned int pPostProcessRGBBuffer_Size = pInferenceModelData->iSrcSize;
|
|
pPostProcessRGBBuffer = new uint8_t[pPostProcessRGBBuffer_Size];
|
|
memcpy(pPostProcessRGBBuffer, pInferenceModelData->pSrcData.get(), pPostProcessRGBBuffer_Size);
|
|
|
|
std::shared_ptr<FrameData> pPostProcessRGBFrameData = std::make_shared<FrameData>();
|
|
|
|
pPostProcessRGBFrameData->iDataSource = engineId_;
|
|
pPostProcessRGBFrameData->iSize = pPostProcessRGBBuffer_Size;
|
|
pPostProcessRGBFrameData->pData.reset(pPostProcessRGBBuffer, [](void* data){if(data) {delete[] data; data = nullptr;}}); //智能指针管理内存
|
|
// pPostProcessRGBFrameData->pData.reset(pPostProcessRGBBuffer, Deleter); //智能指针管理内存
|
|
pPostProcessRGBFrameData->i64TimeStamp = pInferenceModelData->i64TimeStamp;
|
|
|
|
|
|
#if 1
|
|
//构造画框后的RGB数据压入下一引擎
|
|
iRet = outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pPostProcessRGBFrameData));
|
|
if (iRet != APP_ERR_OK){
|
|
LogError << "push the postprocess rgb frame data failed...";
|
|
// std::cerr<<"push the postprocess rgb frame data failed..."<<std::endl;
|
|
}else{
|
|
// std::cout<<"push the postprocess rgb frame data success!"<<std::endl;
|
|
}
|
|
#endif
|
|
}
|
|
}
|
|
|
|
|
|
|