SOMS/test/XuAilibTest/SegDetector.h
2024-07-15 10:31:26 +08:00

58 lines
2.0 KiB
C++

#pragma once
#include "IObjectDetector.h"
#include "dnninfer.h"
#include "opencv2/core/base.hpp"
#include <vector>
class CSegDetector : public IObjectDetector
{
public:
CSegDetector(const char* strModelFile, int iBatchSize = 1, int iInferMode = 1, int iDeviceIndex = 0);
virtual ~CSegDetector();
//初始化对象
virtual bool Init();
//释放对象
virtual bool UnInit();
//获取对象名
virtual const char* GetType() { return "CSegDetector"; };
//执行infer
virtual bool Execute(cv::Mat& image);
virtual bool Execute(std::vector<cv::Mat>& images);
virtual bool ExecuteV2(cv::Mat& image);
virtual bool ExecuteV3(cv::Mat& image);
//获取结果
virtual bool GetResults(std::vector<DetectorResult>& vecDetectorResults);
void EnableMask(bool bEnableMask);
void EnableContourArea(bool bEnableContourArea);
void SetMinAreaAndHeight(int minArea, int minHeight);
protected:
bool ProcessInput(cv::Mat& cvInput, Datum& input, double& fx, double& fy, double& xoffset, double& yoffset);
bool ProcessInputs(std::vector<cv::Mat>& cvInput, Datum& input, double& fx, double& fy, double& xoffset, double& yoffset);
bool ProcessInput(cv::Mat& cvInputImage, cv::Mat& cvProcessedImage, double& fx, double& fy, double& xoffset, double& yoffset);
bool MaskToDetectResult(const cv::Mat& mask, float min_area, float min_height, const cv::Size& image_size, std::vector<DetectorResult>& vecDetectResult);
bool ProbToDetectResult(Datum& probDatum, float min_area, float min_height, const cv::Size& image_size, std::vector<DetectorResult>& vecDetectResult);
private:
bool m_bInit;
int m_iBatchSize;
int m_iInferMode;
int m_iDeviceInedx;
std::string m_strModelFile;
std::vector<DetectorResult> m_vecDetectResult;
void* m_pInfer;
int m_iChannel;
int m_iHeight;
int m_iWidth;
int m_iLabelNum;
int m_iMinArea;
int m_iMinHeight;
Datum m_inputDatum;
Datum m_outputDatum;
Datum m_outputProbDatum;
bool m_bEnableMask;
bool m_bEnableContourArea;
};