106 lines
2.5 KiB
C++
106 lines
2.5 KiB
C++
#include "NanoDetector.h"
|
||
#include "dnninfer.h"
|
||
|
||
CNanoDetector::CNanoDetector(const char* strModelFile, int iBatchSize, int iInferMode, int iDeviceIndex)
|
||
{
|
||
m_bInit = false;
|
||
m_strModelFile = strModelFile;
|
||
m_pInfer = NULL;
|
||
m_iBatchSize = iBatchSize;
|
||
m_iInferMode = iInferMode;
|
||
m_iDeviceInedx = 0;
|
||
}
|
||
CNanoDetector::~CNanoDetector()
|
||
{
|
||
UnInit();
|
||
}
|
||
|
||
bool CNanoDetector::Init()
|
||
{
|
||
bool bRet = false;
|
||
{
|
||
m_pInfer = DnnInfer_Init(m_strModelFile.c_str(), m_iBatchSize, m_iInferMode, m_iDeviceInedx);
|
||
if (m_pInfer != NULL)
|
||
{
|
||
m_bInit = true;
|
||
bRet = true;
|
||
}
|
||
}
|
||
return bRet;
|
||
}
|
||
|
||
bool CNanoDetector::UnInit()
|
||
{
|
||
if (m_bInit)
|
||
{
|
||
m_bInit = false;
|
||
if (m_pInfer != NULL)
|
||
{
|
||
DnnInfer_Close(m_pInfer);
|
||
}
|
||
}
|
||
return true;
|
||
}
|
||
|
||
const char* CNanoDetector::GetType()
|
||
{
|
||
return "CNanoDetector";
|
||
}
|
||
|
||
bool CNanoDetector::Execute(cv::Mat& inputImage)
|
||
{
|
||
m_vecDetectResult.clear();
|
||
if (m_pInfer == NULL)
|
||
return false;
|
||
if (!inputImage.isContinuous())
|
||
inputImage = inputImage.clone();
|
||
Datum inputData(inputImage.data, Shape(1, inputImage.channels(), inputImage.rows, inputImage.cols), 1);
|
||
|
||
DetResult detResult[512];
|
||
int iUseCount = 0;
|
||
|
||
DnnInfer_Det_InferV2(m_pInfer, &inputData, detResult, 512, iUseCount, 0.4, 0.5);
|
||
|
||
for (int i = 0; i < iUseCount; i++)
|
||
{
|
||
DetectorResult detectResult;
|
||
detectResult.hasMask = false;
|
||
detectResult.image_id = 0;
|
||
detectResult.label = detResult[i].label;
|
||
detectResult.score = detResult[i].score;
|
||
strcpy_s(detectResult.name, detResult[i].name);
|
||
char strLabelName[32] = { 0 };
|
||
DnnInfer_GetLabelName(m_pInfer, strLabelName, detResult[i].label);
|
||
//printf("LabelName:%s %d %d %d %d\n", strLabelName);
|
||
//detectResult.name = detResult[i].name;
|
||
detectResult.left = detResult[i].left;
|
||
detectResult.top = detResult[i].top;
|
||
detectResult.width = detResult[i].width;
|
||
detectResult.height = detResult[i].height;
|
||
m_vecDetectResult.emplace_back(detectResult);
|
||
}
|
||
|
||
//cv::Mat mask, predict;
|
||
//int target_label = 1;
|
||
////ʹmask = target_labelµÄÏñËØ = 1
|
||
//cv::Mat target = mask == target_label;
|
||
////ʹ predict = target_labelµÄÏñËØ = 1
|
||
//cv::Mat result = predict = target_label;
|
||
//
|
||
////ʹ±ê×¢ºÍÍÆÀí¶¼µÈÓÚtargetlabelµÄÏñËØ=1
|
||
//cv::Mat intersection = target == result;
|
||
////±ê×¢ºÍÍÆÀí·ÇÁãÏñËØÏà¼Ó - ±ê×¢ºÍÍÆÀí¶¼Îª1µÄÏñËØ
|
||
//cv::Mat unionSet = cv::sum(target + result) - intersection;
|
||
////¼ÆËã½»¼¯=1ÏñËصĺÍ/=1µÄÏà»ú²¢¼¯µÄºÍ
|
||
//float iou = cv::sum(intersection)[0] / cv::sum(unionSet)[0];
|
||
|
||
|
||
return true;
|
||
}
|
||
|
||
bool CNanoDetector::GetResults(std::vector<DetectorResult>& vecDetectorResults)
|
||
{
|
||
vecDetectorResults = m_vecDetectResult;
|
||
return true;
|
||
}
|