DogPoseDetectorOnnx.h
1.12 KB
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
#pragma once
#include <fstream>
#include <iostream>
#include <sstream>
#include <vector>
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "AlgorithmResult.h"
#include "buffers.h"
struct InferDeleter
{
template <typename T>
void operator()(T* obj) const
{
delete obj;
}
};
template <typename T>
using SampleUniquePtr = std::unique_ptr<T, InferDeleter>;
class DogPoseDetectorOnnx
{
public:
DogPoseDetectorOnnx();
~DogPoseDetectorOnnx();
bool init();
std::vector<DogPoseResult> detect(unsigned char *pGpuBgb, int src_width, int src_height);
private:
bool m_bUseFP16{true};
int32_t m_dlaCore{ -1 }; //!< Specify the DLA core to run network on.
std::shared_ptr<nvinfer1::IRuntime> mRuntime; //!< The TensorRT runtime used to deserialize the engine
std::shared_ptr<nvinfer1::ICudaEngine> mEngine; //!< The TensorRT engine used to run the network
void** m_data_buffer;
int m_input_node_index;
nvinfer1::Dims m_input_node_dim;
int m_output_node_index;
nvinfer1::Dims m_output_node_dim;
SampleUniquePtr<nvinfer1::IExecutionContext> context;
};