first commit

third_party/trt_yolov8/plugin/yololayer.h

@@ -0,0 +1,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