forked from linghu8812/tensorrt_inference
-
Notifications
You must be signed in to change notification settings - Fork 0
/
MiniFASNet.h
29 lines (25 loc) · 850 Bytes
/
MiniFASNet.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
#ifndef MINIFASNET_TRT_MINIFASNET_H
#define MINIFASNET_TRT_MINIFASNET_H
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
class MiniFASNet
{
public:
MiniFASNet(const std::string &config_file);
~MiniFASNet();
void LoadEngine();
bool InferenceFolder(const std::string &folder_name);
private:
void EngineInference(const std::vector<std::string> &image_list, const int &outSize,void **buffers,
const std::vector<int64_t> &bufferSize, cudaStream_t stream);
std::vector<float> prepareImage(std::vector<cv::Mat> & vec_img);
std::string onnx_file;
std::string engine_file;
int BATCH_SIZE;
int INPUT_CHANNEL;
int IMAGE_WIDTH;
int IMAGE_HEIGHT;
nvinfer1::ICudaEngine *engine = nullptr;
nvinfer1::IExecutionContext *context = nullptr;
};
#endif //MINIFASNET_TRT_MINIFASNET_H