28 #ifndef TRT_INFERENCE_H_
29 #define TRT_INFERENCE_H_
34 #include "NvCaffeParser.h"
35 #include "NvOnnxParser.h"
36 #include <opencv2/objdetect/objdetect.hpp>
37 using namespace nvinfer1;
38 using namespace nvcaffeparser1;
39 using namespace nvonnxparser;
43 #define GOOGLENET_SINGLE_CLASS 0
44 #define GOOGLENET_THREE_CLASS 1
45 #define RESNET_THREE_CLASS 2
55 int getNetWidth()
const;
57 int getNetHeight()
const;
59 uint32_t getBatchSize()
const;
61 int getChannel()
const;
63 int getModelClassCnt()
const;
65 void* getScales()
const;
67 void* getOffsets()
const;
71 void*& getBuffer(
const int& index);
73 float*& getInputBuf();
75 uint32_t getNumTrtInstances()
const;
78 void setMode(
const int& mode);
80 void setBatchSize(
const uint32_t& batchsize);
82 void setDumpResult(
const bool& dump_result);
84 void setTrtProfilerEnabled(
const bool& enable_trt_profiler);
86 int getFilterNum()
const;
87 void setFilterNum(
const unsigned int& filter_num);
91 void setModelIndex(
int modelIndex);
93 void buildTrtContext(
const string& deployfile,
94 const string& modelfile,
bool bUseCPUBuf =
false,
bool isOnnxModel =
false);
97 queue< vector<cv::Rect> >* rectList_queue,
100 void destroyTrtContext(
bool bUseCPUBuf =
false);
110 float *output_cov_buf;
111 float *output_bbox_buf;
114 float helnet_scale[4];
117 IExecutionContext *context;
118 uint32_t *pResultArray;
126 bool enable_trt_profiler;
128 IHostMemory *trtModelStream{
nullptr};
129 vector<string> outputs;
134 uint64_t elapsed_frame_num;
135 uint64_t elapsed_time;
141 Dims3 outputDimsBBOX;
144 size_t outputSizeBBOX;
155 float input_scale[3];
156 float bbox_output_scales[4];
158 } *g_pModelNetAttr, gModelNetAttr[4] = {
185 {-640, -368, 640, 368},
199 {0.0039215697906911373, 0.0039215697906911373, 0.0039215697906911373},
200 {-640, -368, 640, 368},
209 int parseNet(
const string& deployfile);
210 void parseBbox(vector<cv::Rect>* rectList,
int batch_th);
211 void ParseResnet10Bbox(vector<cv::Rect>* rectList,
int batch_th);
212 void allocateMemory(
bool bUseCPUBuf);
213 void releaseMemory(
bool bUseCPUBuf);
214 void caffeToTRTModel(
const string& deployfile,
const string& modelfile);
215 void onnxToTRTModel(
const string& modelfile);