add samples

wangrong

2025-03-19 291afb9c81aae1d15d9e3edcea618d0a2e1fc5a6

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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
#pragma once 
 
#include <mutex> 
#include <memory> 
#include <string> 
#include <vector> 
#include <chrono> 
#include <iostream> 
#include <fstream> 
#include <cuda_runtime.h> 
#include <cuda_runtime_api.h> 
#include <NvInfer.h> 
#include <NvOnnxConfig.h> 
#include <NvInferPlugin.h> 
#include <opencv2/core/core.hpp> 
#include <opencv2/opencv.hpp> 
#include <sys/time.h> 
#include <device_launch_parameters.h> 
 
#include "../util/algorithm_util.h" 
#include "logging.h" 
 
using namespace nvinfer1; 
using namespace nvonnxparser; 
using namespace trt_vehicle; 
 
//extern "C" void calGpuDistance(float** f1,float** f2,float** dis,int outN,int outC); 
 
namespace trt_vehicle{ 
 
#define CheckCudaError(status) \ 
{ \ 
    if (status != 0) \ 
    { \ 
        return; \ 
    } \ 
} 
    enum class NetworkInputType 
    { 
        HWC, 
        CHW 
    }; 
 
class CudaPredictor{ 
public: 
    ~CudaPredictor(); 
public: 
    // create engine from engine path 
    int init(const string path,int deviceID = 0, bool verbose=false); 
 
    //get input/output sizes of model 
    int getSizeYolo(int& batch,int& inputSizeC,int& inputSizeH,int& inputSizeW,int& outputNum,int& classNum,int& boxNum);  
 
    //get input/output sizes of model 
    int getSize(int& batch,int& inputSizeC,int& inputSizeH,int& inputSizeW,int& outputDim1,int& outputDim2,int& outputDim3, NetworkInputType networkInputType = NetworkInputType::HWC); 
 
    // TODO 车辆分类模型 
    int getSizeVehicle(int& batch,int& inputSizeC,int& inputSizeH,int& inputSizeW,int& outputDim1,int& outputDim2,int& outputDim3); 
     
    //forward of neural network 
    int infer(vector<void *> &buffers, int batch=1); 
 
private: 
    int loadModel(const string &path); 
    int prepare(); 
private: 
    nvinfer1::IRuntime* m_runtime = nullptr; 
    nvinfer1::ICudaEngine* m_engine = nullptr; 
    nvinfer1::IExecutionContext* m_context = nullptr; 
    cudaStream_t m_stream = nullptr; 
}; 
 
class BaseModel{ 
    public: 
        BaseModel(NetworkInputType inputType = NetworkInputType::HWC, std::string meanFile = "") 
                { m_inputType = inputType; m_MeanFile = meanFile; } 
 
        virtual ~BaseModel(); 
 
        //input dimensions of the model  
        int getInputSize(int& inputW,int& inputH); 
 
        //image preprocessimg 
        int imagePreprocess(cv::Mat& img,vector<float>& imageData); 
 
        bool readMeanImageFile(); 
    protected: 
        //image resize 
        void imageResize(const cv::Mat &src,cv::Size size,std::vector<float>& data); 
 
    protected: 
        int m_inputSizeH; 
        int m_inputSizeW; 
        int m_inputSizeC; 
        int m_batch; 
        int m_downScale; 
 
        NetworkInputType m_inputType; 
        std::string m_MeanFile; 
        float* m_meanFileData = nullptr; 
    }; 
}