28 #ifndef TRT_INFERENCE_H_
29 #define TRT_INFERENCE_H_
34 #include "NvCaffeParser.h"
35 #include "opencv2/video/tracking.hpp"
36 #include "opencv2/imgproc/imgproc.hpp"
37 #include "opencv2/highgui/highgui.hpp"
38 #include <opencv2/objdetect/objdetect.hpp>
39 using namespace nvinfer1;
40 using namespace nvcaffeparser1;
44 #define GOOGLENET_SINGLE_CLASS 0
45 #define GOOGLENET_THREE_CLASS 1
46 #define RESNET_THREE_CLASS 2
56 int getNetWidth()
const;
58 int getNetHeight()
const;
60 uint32_t getBatchSize()
const;
62 int getChannel()
const;
64 int getModelClassCnt()
const;
66 void* getScales()
const;
68 void* getOffsets()
const;
72 void*& getBuffer(
const int& index);
74 float*& getInputBuf();
76 uint32_t getNumTrtInstances()
const;
79 void setMode(
const int& mode);
81 void setBatchSize(
const uint32_t& batchsize);
83 void setDumpResult(
const bool& dump_result);
85 void setTrtProfilerEnabled(
const bool& enable_trt_profiler);
87 int getFilterNum()
const;
88 void setFilterNum(
const unsigned int& filter_num);
92 void setModelIndex(
int modelIndex);
94 void buildTrtContext(
const string& deployfile,
95 const string& modelfile,
bool bUseCPUBuf =
false);
98 queue< vector<cv::Rect> >* rectList_queue,
101 void destroyTrtContext(
bool bUseCPUBuf =
false);
111 float *output_cov_buf;
112 float *output_bbox_buf;
115 float helnet_scale[4];
118 IExecutionContext *context;
119 uint32_t *pResultArray;
127 bool enable_trt_profiler;
128 IHostMemory *trtModelStream{
nullptr};
129 vector<string> outputs;
134 uint64_t elapsed_frame_num;
135 uint64_t elapsed_time;
141 DimsCHW 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);
const char * INPUT_BLOB_NAME
const char * OUTPUT_BBOX_NAME
const char * OUTPUT_BLOB_NAME