-
Notifications
You must be signed in to change notification settings - Fork 545
/
yolo_layer.h
150 lines (106 loc) · 5.42 KB
/
yolo_layer.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
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
#ifndef _YOLO_LAYER_H
#define _YOLO_LAYER_H
#include <cassert>
#include <vector>
#include <string>
#include <iostream>
#include "math_constants.h"
#include "NvInfer.h"
#define MAX_ANCHORS 6
#if NV_TENSORRT_MAJOR >= 8
#define NOEXCEPT noexcept
#else
#define NOEXCEPT
#endif
#define CHECK(status) \
do { \
auto ret = status; \
if (ret != 0) { \
std::cerr << "Cuda failure in file '" << __FILE__ \
<< "' line " << __LINE__ \
<< ": " << ret << std::endl; \
abort(); \
} \
} while (0)
namespace Yolo
{
static constexpr float IGNORE_THRESH = 0.01f;
struct alignas(float) Detection {
float bbox[4]; // x, y, w, h
float det_confidence;
float class_id;
float class_confidence;
};
}
namespace nvinfer1
{
class YoloLayerPlugin: public IPluginV2IOExt
{
public:
YoloLayerPlugin(int yolo_width, int yolo_height, int num_anchors, float* anchors, int num_classes, int input_width, int input_height, float scale_x_y, int new_coords);
YoloLayerPlugin(const void* data, size_t length);
~YoloLayerPlugin() override = default;
IPluginV2IOExt* clone() const NOEXCEPT override;
int initialize() NOEXCEPT override { return 0; }
void terminate() NOEXCEPT override;
void destroy() NOEXCEPT override { delete this; }
size_t getSerializationSize() const NOEXCEPT override;
void serialize(void* buffer) const NOEXCEPT override;
int getNbOutputs() const NOEXCEPT override { return 1; }
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override;
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0; }
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT override { return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; }
const char* getPluginType() const NOEXCEPT override { return "YoloLayer_TRT"; }
const char* getPluginVersion() const NOEXCEPT override { return "1"; }
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override { mPluginNamespace = pluginNamespace; }
const char* getPluginNamespace() const NOEXCEPT override { return mPluginNamespace; }
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const NOEXCEPT override { return DataType::kFLOAT; }
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override { return false; }
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override {}
//using IPluginV2IOExt::configurePlugin;
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) NOEXCEPT override {}
void detachFromContext() NOEXCEPT override {}
#if NV_TENSORRT_MAJOR >= 8
int32_t enqueue(int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override;
#else
int enqueue(int batchSize, const void* const * inputs, void** outputs, void* workspace, cudaStream_t stream) NOEXCEPT override;
#endif
private:
void forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int batchSize = 1);
int mThreadCount = 64;
int mYoloWidth, mYoloHeight, mNumAnchors;
float mAnchorsHost[MAX_ANCHORS * 2];
float *mAnchors; // allocated on GPU
int mNumClasses;
int mInputWidth, mInputHeight;
float mScaleXY;
int mNewCoords = 0;
const char* mPluginNamespace;
};
class YoloPluginCreator : public IPluginCreator
{
public:
YoloPluginCreator();
~YoloPluginCreator() override = default;
const char* getPluginName() const NOEXCEPT override;
const char* getPluginVersion() const NOEXCEPT override;
const PluginFieldCollection* getFieldNames() NOEXCEPT override;
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) NOEXCEPT override;
IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) NOEXCEPT override;
void setPluginNamespace(const char* libNamespace) NOEXCEPT override
{
mNamespace = libNamespace;
}
const char* getPluginNamespace() const NOEXCEPT override
{
return mNamespace.c_str();
}
private:
static PluginFieldCollection mFC;
static std::vector<PluginField> mPluginAttributes;
std::string mNamespace;
};
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
};
#endif