initial commit

wangrong

2025-01-03 2c4b8838f89bebca97615c5f0bac2c58e0349771

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
#pragma once 
#include <string> 
#include <vector> 
#include "NvInfer.h" 
#include "../include/macros.h" 
namespace nvinfer1 { 
class API YoloLayerPlugin : public IPluginV2IOExt { 
   public: 
    YoloLayerPlugin(int classCount, int numberofpoints, float confthreshkeypoints, int netWidth, int netHeight, 
                    int maxOut, bool is_segmentation, bool is_pose, const int* strides, int stridesLength); 
 
    YoloLayerPlugin(const void* data, size_t length); 
    ~YoloLayerPlugin(); 
 
    int getNbOutputs() const TRT_NOEXCEPT override { return 1; } 
 
    nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; 
 
    int initialize() TRT_NOEXCEPT override; 
 
    virtual void terminate() TRT_NOEXCEPT override {} 
 
    virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } 
 
    virtual int enqueue(int batchSize, const void* const* inputs, void* TRT_CONST_ENQUEUE* outputs, void* workspace, 
                        cudaStream_t stream) TRT_NOEXCEPT override; 
 
    virtual size_t getSerializationSize() const TRT_NOEXCEPT override; 
 
    virtual void serialize(void* buffer) const TRT_NOEXCEPT override; 
 
    bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, 
                                   int nbOutputs) const TRT_NOEXCEPT override { 
        return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; 
    } 
 
    const char* getPluginType() const TRT_NOEXCEPT override; 
 
    const char* getPluginVersion() const TRT_NOEXCEPT override; 
 
    void destroy() TRT_NOEXCEPT override; 
 
    IPluginV2IOExt* clone() const TRT_NOEXCEPT override; 
 
    void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; 
 
    const char* getPluginNamespace() const TRT_NOEXCEPT override; 
 
    nvinfer1::DataType getOutputDataType(int32_t index, nvinfer1::DataType const* inputTypes, 
                                         int32_t nbInputs) const TRT_NOEXCEPT; 
 
    bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, 
                                      int nbInputs) const TRT_NOEXCEPT override; 
 
    bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; 
 
    void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, 
                         IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; 
 
    void configurePlugin(PluginTensorDesc const* in, int32_t nbInput, PluginTensorDesc const* out, 
                         int32_t nbOutput) TRT_NOEXCEPT override; 
 
    void detachFromContext() TRT_NOEXCEPT override; 
 
   private: 
    void forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight, 
                    int mYoloV8NetWidth, int batchSize); 
    int mThreadCount = 256; 
    const char* mPluginNamespace; 
    int mClassCount; 
    int mNumberofpoints; 
    float mConfthreshkeypoints; 
    int mYoloV8NetWidth; 
    int mYoloV8netHeight; 
    int mMaxOutObject; 
    bool is_segmentation_; 
    bool is_pose_; 
    int* mStrides; 
    int mStridesLength; 
}; 
 
class API YoloPluginCreator : public IPluginCreator { 
   public: 
    YoloPluginCreator(); 
    ~YoloPluginCreator() override = default; 
 
    const char* getPluginName() const TRT_NOEXCEPT override; 
 
    const char* getPluginVersion() const TRT_NOEXCEPT override; 
 
    const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; 
 
    nvinfer1::IPluginV2IOExt* createPlugin(const char* name, 
                                           const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override; 
 
    nvinfer1::IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, 
                                                size_t serialLength) TRT_NOEXCEPT override; 
 
    void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override { mNamespace = libNamespace; } 
 
    const char* getPluginNamespace() const TRT_NOEXCEPT override { return mNamespace.c_str(); } 
 
   private: 
    std::string mNamespace; 
    static PluginFieldCollection mFC; 
    static std::vector<PluginField> mPluginAttributes; 
}; 
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); 
}  // namespace nvinfer1