SOMS/test/XuAilibTest/SegDetector.cpp

649 lines
18 KiB
C++
Raw Normal View History

2024-07-15 10:31:26 +08:00
#include "SegDetector.h"
#include "DnnDefine.h"
#ifndef _myminmax
#define _myminmax
#define mymin(a,b) (a) < (b) ? (a) : (b)
#define mymax(a,b) (a) > (b) ? (a) : (b)
#endif
CSegDetector::CSegDetector(const char* strModelFile, int iBatchSize, int iInferMode, int iDeviceIndex)
{
{
static bool bFirst = true;
if (bFirst)
{
DnnInfer_InitGlobalLog();
bFirst = false;
}
}
m_bInit = false;
m_strModelFile = strModelFile;
m_pInfer = NULL;
m_iMinArea = 20;
m_iMinHeight = 5;
m_bEnableMask = true;
m_bEnableContourArea = true;
m_iBatchSize = iBatchSize;
m_iInferMode = iInferMode;
m_iDeviceInedx = iDeviceIndex;
}
CSegDetector::~CSegDetector()
{
UnInit();
}
bool CSegDetector::Init()
{
bool bRet = false;
{
InferenceParameterEx inferParam;
DnnInfer_GetInferParamFormModelFile(m_strModelFile.c_str(), &inferParam);
m_pInfer = DnnInfer_Init(m_strModelFile.c_str(), m_iBatchSize, m_iInferMode, m_iDeviceInedx);
if (m_pInfer != NULL)
{
m_bInit = true;
DnnInfer_GetInputInfo(m_pInfer, &m_iWidth, &m_iHeight, &m_iChannel, &m_iBatchSize, &m_iLabelNum);
printf("GetInputInfo %d %d %d\n", m_iWidth, m_iHeight, m_iChannel);
GpuDeviceInfo gpuDeviceInfo;
DnnInfer_GetDeviceInfo(&gpuDeviceInfo);
for (int i = 0; i < gpuDeviceInfo.gpu_count; i++)
{
std::cout << "<EFBFBD>Կ<EFBFBD><EFBFBD><EFBFBD><EFBFBD>ƣ<EFBFBD>" << gpuDeviceInfo.gpu_name[i] << std::endl;
std::cout << "<EFBFBD>Դ<EFBFBD><EFBFBD><EFBFBD>С<EFBFBD><EFBFBD>" << gpuDeviceInfo.gpu_memory_size[i] << " MB" << std::endl;
std::cout << "һ<EFBFBD><EFBFBD>block<EFBFBD>Ĺ<EFBFBD><EFBFBD><EFBFBD><EFBFBD>ڴ<EFBFBD><EFBFBD><EFBFBD>С<EFBFBD><EFBFBD>" << gpuDeviceInfo.gpu_sharedMemPerBlock[i] << " KB" << std::endl;
std::cout << "block<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>߳<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>" << gpuDeviceInfo.gpu_maxThreadsPerBlock[i] << std::endl;
}
//m_iMinArea = 20;
//m_iMinHeight = 5;
bRet = true;
}
}
return bRet;
}
bool CSegDetector::UnInit()
{
if (m_bInit)
{
m_bInit = false;
if (m_pInfer != NULL)
{
DnnInfer_Close(m_pInfer);
}
}
return true;
}
bool CSegDetector::Execute(cv::Mat& image)
{
if (!m_bInit)
return false;
m_vecDetectResult.clear();
m_iHeight = image.rows;
m_iWidth = image.cols;
double dFx, dFy, dXoffset, dYoffset;
ProcessInput(image, m_inputDatum, dFx, dFy, dXoffset, dYoffset);
std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
//Infer(pObject, &input, iLabelArray, fpProbalityArray);
DnnInfer_Seg_Infer(m_pInfer, &m_inputDatum, &m_outputDatum);
auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
int iMaskSize = m_iHeight * m_iWidth;
cv::Mat mask(m_iHeight, m_iWidth, CV_8UC1, (uchar*)((uchar*)m_outputDatum.data));
std::vector<DetectorResult> vecDetectResult;
MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);
for (int j = 0; j < vecDetectResult.size(); j++)
{
DetectorResult detectorResult = vecDetectResult[j];
detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
detectorResult.width = mymin(vecDetectResult[j].width / dFx, image.cols - detectorResult.left);
detectorResult.height = mymin(vecDetectResult[j].height / dFy, image.rows - detectorResult.top);
if (detectorResult.width > 0 && detectorResult.height > 10)
{
if (m_bEnableMask)
{
if (!detectorResult.mask.empty())
cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
}
if (m_bEnableContourArea)
{
detectorResult.area *= dFx * dFy;
}
m_vecDetectResult.emplace_back(detectorResult);
}
}
return true;
}
bool CSegDetector::Execute(std::vector<cv::Mat>& images)
{
if (!m_bInit)
return false;
m_vecDetectResult.clear();
m_iHeight = images[0].rows;
m_iWidth = images[0].cols;
double dFx, dFy, dXoffset, dYoffset;
ProcessInputs(images, m_inputDatum, dFx, dFy, dXoffset, dYoffset);
std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
//Infer(pObject, &input, iLabelArray, fpProbalityArray);
DnnInfer_Seg_Infer(m_pInfer, &m_inputDatum, &m_outputDatum);
for (int i = 0; i < m_outputDatum.shape.N; i++)
{
auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
int iMaskSize = m_iHeight * m_iWidth;
cv::Mat mask(m_iHeight, m_iWidth, CV_8UC1, (uchar*)((uchar*)m_outputDatum.data + i * iMaskSize));
std::vector<DetectorResult> vecDetectResult;
MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);
for (int j = 0; j < vecDetectResult.size(); j++)
{
DetectorResult detectorResult = vecDetectResult[j];
detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
detectorResult.width = mymin(vecDetectResult[j].width / dFx, m_iWidth - detectorResult.left);
detectorResult.height = mymin(vecDetectResult[j].height / dFy, m_iHeight - detectorResult.top);
if (detectorResult.width > 0 && detectorResult.height > 10)
{
if (m_bEnableMask)
{
if (!detectorResult.mask.empty())
cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
}
if (m_bEnableContourArea)
{
detectorResult.area *= dFx * dFy;
}
m_vecDetectResult.emplace_back(detectorResult);
}
}
}
return true;
}
bool CSegDetector::ExecuteV2(cv::Mat& image)
{
if (!m_bInit)
return false;
m_vecDetectResult.clear();
m_iHeight = image.rows;
m_iWidth = image.cols;
double dFx, dFy, dXoffset, dYoffset;
ProcessInput(image, m_inputDatum, dFx, dFy, dXoffset, dYoffset);
std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
//Infer(pObject, &input, iLabelArray, fpProbalityArray);
DnnInfer_Seg_InferV2(m_pInfer, &m_inputDatum, &m_outputDatum, &m_outputProbDatum);
auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
int iMaskSize = m_iHeight * m_iWidth;
cv::Mat mask(m_iHeight, m_iWidth, CV_8UC1, (uchar*)((uchar*)m_outputDatum.data));
cv::Mat maxProbMask(m_iHeight, m_iWidth, CV_32FC1, ((float*)m_outputProbDatum.data));
cv::imshow("maxProb", maxProbMask);
cv::waitKey(0);
std::vector<DetectorResult> vecDetectResult;
//ProbToDetectResult(m_outputProbDatum, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);
MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);
for (int j = 0; j < vecDetectResult.size(); j++)
{
DetectorResult detectorResult = vecDetectResult[j];
detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
detectorResult.width = mymin(vecDetectResult[j].width / dFx, image.cols - detectorResult.left);
detectorResult.height = mymin(vecDetectResult[j].height / dFy, image.rows - detectorResult.top);
if (detectorResult.width > 0 && detectorResult.height > 10)
{
if (m_bEnableMask)
{
if (!detectorResult.mask.empty())
{
cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
}
}
if (m_bEnableContourArea)
{
detectorResult.area *= dFx * dFy;
}
m_vecDetectResult.emplace_back(detectorResult);
}
}
return true;
}
bool CSegDetector::ExecuteV3(cv::Mat& image)
{
if (!m_bInit)
return false;
m_vecDetectResult.clear();
m_iHeight = image.rows;
m_iWidth = image.cols;
double dFx, dFy, dXoffset, dYoffset;
ProcessInput(image, m_inputDatum, dFx, dFy, dXoffset, dYoffset);
std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
//Infer(pObject, &input, iLabelArray, fpProbalityArray);
SegResult segResult;
memset(&segResult, 0, sizeof(segResult));
DnnInfer_Seg_InferV3(m_pInfer, &m_inputDatum, &segResult);
auto elapsedSeconds = std::chrono::steady_clock::now() - startTime;
printf("infer cost time %lf\n", std::chrono::duration<double, std::milli>(elapsedSeconds).count());
int iMaskSize = m_iHeight * m_iWidth;
cv::Mat mask(segResult.inferShape.H, segResult.inferShape.W, CV_8UC1, (segResult.pLabel));
cv::Mat maxProbMask(segResult.inferShape.H, segResult.inferShape.W, CV_32FC1, (segResult.pProb + 1 * segResult.inferShape.H * segResult.inferShape.W));
cv::imshow("maxProb", maxProbMask);
cv::waitKey(0);
std::vector<DetectorResult> vecDetectResult;
dFx = segResult.dScale;
dFy = segResult.dScale;
dXoffset = segResult.offsetX;
dYoffset = segResult.offsetY;
//ProbToDetectResult(m_outputProbDatum, m_iMinArea, m_iMinHeight, cv::Size(m_iWidth, m_iHeight), vecDetectResult);
MaskToDetectResult(mask, m_iMinArea, m_iMinHeight, cv::Size(segResult.inferShape.W, segResult.inferShape.H), vecDetectResult);
for (int j = 0; j < vecDetectResult.size(); j++)
{
DetectorResult detectorResult = vecDetectResult[j];
detectorResult.left = mymax((vecDetectResult[j].left - dXoffset) / dFx, 0.0);
detectorResult.top = mymax((vecDetectResult[j].top - dYoffset) / dFy, 0.0);
detectorResult.width = mymin(vecDetectResult[j].width / dFx, image.cols - detectorResult.left);
detectorResult.height = mymin(vecDetectResult[j].height / dFy, image.rows - detectorResult.top);
if (detectorResult.width > 0 && detectorResult.height > 10)
{
if (m_bEnableMask)
{
if (!detectorResult.mask.empty())
{
cv::resize(detectorResult.mask, detectorResult.mask, cv::Size(detectorResult.width, detectorResult.height), cv::INTER_NEAREST);
}
}
if (m_bEnableContourArea)
{
detectorResult.area *= dFx * dFy;
}
m_vecDetectResult.emplace_back(detectorResult);
}
}
return true;
}
bool CSegDetector::GetResults(std::vector<DetectorResult>& vecDetectorResults)
{
vecDetectorResults = m_vecDetectResult;
return true;
}
void CSegDetector::EnableMask(bool bEnableMask)
{
m_bEnableMask = bEnableMask;
}
void CSegDetector::EnableContourArea(bool bEnableContourArea)
{
m_bEnableContourArea = bEnableContourArea;
}
void CSegDetector::SetMinAreaAndHeight(int minArea, int minHeight)
{
m_iMinArea = minArea;
m_iMinHeight = minHeight;
}
bool CSegDetector::ProcessInput(cv::Mat& cvInput, Datum& input, double& dFx, double& dFy, double& dXoffset, double& dYoffset)
{
cv::Mat resizeImage;
int iModelWidth = m_iWidth;
int iModelHeight = m_iHeight;
float fResizeScale = 1.0;
int padoffsetX = 0;
int padoffsetY = 0;
dFx = 1.0;
dFy = 1.0;
dXoffset = 0;
dYoffset = 0;
Shape shape(1, cvInput.channels(), cvInput.rows, cvInput.cols);
input.Reshape(shape);
cv::Size input_size = cvInput.size();
const int channels = cvInput.channels();
input.SetDataItemSize(1);
uchar* datum_data = (uchar*)input.Getdata();
unsigned int offsetN = shape.C * shape.H * shape.W;
if (cvInput.isContinuous())
{
memcpy(datum_data, cvInput.data, offsetN);
}
return true;
if (false && (cvInput.rows != iModelHeight || cvInput.cols != iModelWidth))
{
int iheight = cvInput.rows;
int iwidth = cvInput.cols;
double dHeightResizeScale = 1.0 * iModelHeight / iheight;
double dWidthResizeScale = 1.0 * iModelWidth / iwidth;
if (dHeightResizeScale < dWidthResizeScale)
{
fResizeScale = dHeightResizeScale;
}
else
{
fResizeScale = dWidthResizeScale;
}
cv::resize(cvInput, resizeImage, cv::Size(), fResizeScale, fResizeScale);
}
else
{
resizeImage = cvInput;
}
if (m_iChannel == 1)
{
if (resizeImage.channels() == 3)
{
cv::cvtColor(resizeImage, resizeImage, cv::COLOR_BGR2GRAY);
}
cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC1);
padoffsetX = (iModelWidth - resizeImage.cols) / 2;
padoffsetY = (iModelHeight - resizeImage.rows) / 2;
resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));
//cvProcessedImage = paddMat;
Shape shape(1, paddMat.channels(), paddMat.rows, paddMat.cols);
input.Reshape(shape);
cv::Size input_size = paddMat.size();
const int channels = paddMat.channels();
input.SetDataItemSize(1);
uchar* datum_data = (uchar*)input.Getdata();
unsigned int offsetN = shape.C * shape.H * shape.W;
if (paddMat.isContinuous())
{
memcpy(datum_data, paddMat.data, offsetN);
}
}
else
{
if (resizeImage.channels() == 1)
{
cv::cvtColor(resizeImage, resizeImage, cv::COLOR_GRAY2BGR);
}
cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC3);
padoffsetX = (iModelWidth - resizeImage.cols) / 2;
padoffsetY = (iModelHeight - resizeImage.rows) / 2;
resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));
Shape shape(1, paddMat.channels(), paddMat.rows, paddMat.cols);
input.Reshape(shape);
cv::Size input_size = paddMat.size();
const int channels = paddMat.channels();
input.SetDataItemSize(1);
uchar* datum_data = (uchar*)input.Getdata();
unsigned int offsetN = shape.C * shape.H * shape.W;
if (paddMat.isContinuous())
{
memcpy(datum_data, paddMat.data, offsetN);
}
}
dFx = fResizeScale;
dFy = fResizeScale;
dXoffset = padoffsetX;
dYoffset = padoffsetY;
return true;
}
bool CSegDetector::ProcessInputs(std::vector<cv::Mat>& cvInput, Datum& input, double& dFx, double& dFy, double& dXoffset, double& dYoffset)
{
cv::Mat resizeImage;
int iModelWidth = m_iWidth;
int iModelHeight = m_iHeight;
float fResizeScale = 1.0;
int padoffsetX = 0;
int padoffsetY = 0;
dFx = 1.0;
dFy = 1.0;
Shape shape(cvInput.size(), cvInput[0].channels(), cvInput[0].rows, cvInput[0].cols);
input.SetDataItemSize(1);
input.Reshape(shape);
uchar* datum_data = (uchar*)input.Getdata();
for (int i = 0; i < cvInput.size(); i++)
{
unsigned int offsetN = shape.C * shape.H * shape.W;
if (cvInput[0].isContinuous())
{
memcpy(datum_data + i * offsetN, cvInput[0].data, offsetN);
}
}
return true;
}
bool CSegDetector::ProcessInput(cv::Mat& cvInputImage, cv::Mat& cvProcessedImage, double& dFx, double& dFy, double& dXoffset, double& dYoffset)
{
cv::Mat resizeImage;
int iModelWidth = m_iWidth;
int iModelHeight = m_iHeight;
float fResizeScale = 1.0;
int padoffsetX = 0;
int padoffsetY = 0;
if (true || cvInputImage.rows != iModelHeight || cvInputImage.cols != iModelWidth)
{
int iheight = cvInputImage.rows;
int iwidth = cvInputImage.cols;
double dHeightResizeScale = 1.0 * iModelHeight / iheight;
double dWidthResizeScale = 1.0 * iModelWidth / iwidth;
if (dHeightResizeScale < dWidthResizeScale)
{
fResizeScale = dHeightResizeScale;
}
else
{
fResizeScale = dWidthResizeScale;
}
cv::resize(cvInputImage, resizeImage, cv::Size(), fResizeScale, fResizeScale);
}
else
{
resizeImage = cvInputImage;
}
if (m_iChannel == 1)
{
if (resizeImage.channels() == 3)
{
cv::cvtColor(resizeImage, resizeImage, cv::COLOR_BGR2GRAY);
}
cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC1);
padoffsetX = (iModelWidth - resizeImage.cols) / 2;
padoffsetY = (iModelHeight - resizeImage.rows) / 2;
resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));
cvProcessedImage = paddMat;
}
else
{
cv::Mat paddMat = cv::Mat::zeros(cv::Size(iModelWidth, iModelHeight), CV_8UC3);
padoffsetX = (iModelWidth - resizeImage.cols) / 2;
padoffsetY = (iModelHeight - resizeImage.rows) / 2;
resizeImage.copyTo(paddMat(cv::Rect(padoffsetX, padoffsetY, resizeImage.cols, resizeImage.rows)));
cvProcessedImage = paddMat;
}
dFx = fResizeScale;
dFy = fResizeScale;
dXoffset = padoffsetX;
dYoffset = padoffsetY;
return true;
}
bool CSegDetector::MaskToDetectResult(const cv::Mat& mask, float min_area, float min_height,
const cv::Size& image_size, std::vector<DetectorResult>& vecDetectResult) {
double min_val;
double max_val;
cv::minMaxLoc(mask, &min_val, &max_val);
int max_bbox_idx = static_cast<int>(max_val);
cv::Mat resized_mask;
if (mask.size() != image_size)
{
cv::resize(mask, resized_mask, image_size, 0, 0, cv::INTER_NEAREST);
}
else
{
resized_mask = mask.clone();
}
for (int i = 1; i <= max_bbox_idx; i++) {
cv::Mat bbox_mask = resized_mask == i;
std::vector<std::vector<cv::Point>> contours;
cv::findContours(bbox_mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
if (contours.empty())
continue;
for (int j = 0; j < contours.size(); j++)
{
cv::Rect rect = cv::boundingRect(contours[j]);
if (std::min(rect.size().width, rect.size().height) < min_height)
continue;
if (rect.size().area() < min_area)
continue;
DetectorResult detectResult;
detectResult.left = rect.x;
detectResult.top = rect.y;
detectResult.width = rect.width;
detectResult.height = rect.height;
detectResult.label = i;
detectResult.score = 1.0;
if (m_bEnableContourArea)
{
detectResult.area = cv::contourArea(contours[j]);
}
else
{
detectResult.area = detectResult.width * detectResult.height;
}
//detectResult.area = cv::contourArea(contours[j]);
sprintf_s(detectResult.name, "%d", detectResult.label);
if (m_bEnableMask)
{
detectResult.hasMask = true;
detectResult.mask = mask(rect).clone();
}
else
{
detectResult.hasMask = false;
}
vecDetectResult.emplace_back(detectResult);
}
}
return true;
}
bool CSegDetector::ProbToDetectResult(Datum& probDatum, float min_area, float min_height, const cv::Size& image_size, std::vector<DetectorResult>& vecDetectResult)
{
int score_map_size = probDatum.shape.C * probDatum.shape.H * probDatum.shape.W;
const float* score_map_data =
reinterpret_cast<const float*>(probDatum.data);
int num_map_pixels = probDatum.shape.H * probDatum.shape.W;
for (int i = 0; i < probDatum.shape.N; ++i) {
const float* current_start_ptr = score_map_data + i * score_map_size;
cv::Mat ori_score_mat(probDatum.shape.C,
probDatum.shape.H * probDatum.shape.W,
CV_32FC1, const_cast<float*>(current_start_ptr));
ori_score_mat = ori_score_mat.t();
cv::Mat score_mat(probDatum.shape.H, probDatum.shape.W, CV_32FC1);
cv::Mat label_mat(probDatum.shape.H, probDatum.shape.W, CV_8UC1);
for (int j = 0; j < ori_score_mat.rows; ++j) {
double max_value;
cv::Point max_id;
minMaxLoc(ori_score_mat.row(j), 0, &max_value, 0, &max_id);
if (max_value > 0.5)
{
score_mat.at<float>(j) = max_value;
label_mat.at<uchar>(j) = max_id.x;
}
else
{
score_mat.at<float>(j) = max_value;
label_mat.at<uchar>(j) = 0;
}
}
MaskToDetectResult(label_mat, min_area, min_height, image_size, vecDetectResult);
}
return true;
}