111 lines
4.2 KiB
C++
111 lines
4.2 KiB
C++
#pragma once
|
|
#include <string>
|
|
#include <vector>
|
|
#include "NvInfer.h"
|
|
#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, bool is_obb, 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_;
|
|
bool is_obb_;
|
|
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
|