51 lines
2.0 KiB
C
51 lines
2.0 KiB
C
|
#ifndef RETINAFACE_H
|
||
|
|
#define RETINAFACE_H
|
||
|
|
#pragma once
|
||
|
|
#include "ANSFRCommon.h"
|
||
|
|
#define CLIP(a, b, c) (MAX(MIN(a, c), b)) // MIN, MAX defined in opencv
|
||
|
|
namespace ANSCENTER {
|
||
|
|
// Retina will be FD subclass
|
||
|
|
class RetinaFaceTRT {
|
||
|
|
public:
|
||
|
|
RetinaFaceTRT();
|
||
|
|
~RetinaFaceTRT();
|
||
|
|
std::vector<struct Bbox> FindFace(cv::Mat& img);
|
||
|
|
bool Initialize(const std::string engineFile,
|
||
|
|
int frameWidth,
|
||
|
|
int frameHeight,
|
||
|
|
std::string inputName,
|
||
|
|
std::vector<std::string> outputNames,
|
||
|
|
std::vector<int> inputShape,
|
||
|
|
int maxBatchSize,
|
||
|
|
int maxFacesPerScene,
|
||
|
|
float nms_threshold,
|
||
|
|
float bbox_threshold);
|
||
|
|
|
||
|
|
private:
|
||
|
|
ANSCENTER::SPDLogger& _logger = ANSCENTER::SPDLogger::GetInstance("ANSFD", true);
|
||
|
|
TRTLogger m_logger;
|
||
|
|
|
||
|
|
int m_frameWidth, m_frameHeight, m_INPUT_C, m_INPUT_H, m_INPUT_W, m_INPUT_SIZE, m_OUTPUT_SIZE_BASE, m_maxBatchSize, m_maxFacesPerScene;
|
||
|
|
float m_nms_threshold, m_bbox_threshold, m_scale_h, m_scale_w;
|
||
|
|
cv::Mat m_input;
|
||
|
|
float* m_output0, * m_output1;
|
||
|
|
std::vector<struct Bbox> m_outputBbox;
|
||
|
|
|
||
|
|
nvinfer1::ICudaEngine* m_engine;
|
||
|
|
nvinfer1::IExecutionContext* m_context;
|
||
|
|
cudaStream_t stream;
|
||
|
|
void* buffers[3];
|
||
|
|
int inputIndex, outputIndex0, outputIndex1;
|
||
|
|
private:
|
||
|
|
void LoadEngine(const std::string engineFile);
|
||
|
|
void PreInference(std::string inputNames, std::vector<std::string> outputNames);
|
||
|
|
void RunInference(float* input, float* output0, float* output1);
|
||
|
|
void PreProcess(cv::Mat& img);
|
||
|
|
void PostProcessing(float* bbox, float* conf);
|
||
|
|
void CreateAnchorRetinaface(std::vector<AnchorBox>& anchor, int w, int h);
|
||
|
|
static inline bool MCMP(Bbox a, Bbox b);
|
||
|
|
void NMS(std::vector<Bbox>& input_boxes, float NMS_THRESH);
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
#endif
|