#ifndef _DETECT_H_
|
#define _DETECT_H_
|
|
#include <string>
|
#include <vector>
|
#include "NvInfer.h"
|
|
namespace nvinfer1
|
{
|
template <typename T>
|
void write(char*& buffer, const T& val)
|
{
|
*reinterpret_cast<T*>(buffer) = val;
|
buffer += sizeof(T);
|
}
|
|
template <typename T>
|
void read(const char*& buffer, T& val)
|
{
|
val = *reinterpret_cast<const T*>(buffer);
|
buffer += sizeof(T);
|
}
|
|
class Detect :public IPluginV2
|
{
|
public:
|
Detect();
|
Detect(const void* data, size_t length);
|
Detect(const uint32_t n_anchor_, const uint32_t _n_classes_,
|
const uint32_t n_grid_h_, const uint32_t n_grid_w_/*,
|
const uint32_t &n_stride_h_, const uint32_t &n_stride_w_*/);
|
~Detect();
|
int getNbOutputs()const noexcept override
|
{
|
return 1;
|
}
|
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) noexcept override
|
{
|
return inputs[0];
|
}
|
int initialize() noexcept override
|
{
|
return 0;
|
}
|
void terminate() noexcept override
|
{
|
}
|
size_t getWorkspaceSize(int maxBatchSize) const noexcept override
|
{
|
return 0;
|
}
|
int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace,
|
cudaStream_t stream) noexcept override;
|
|
|
bool supportsFormat(DataType type, PluginFormat format) const noexcept override;
|
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept override;
|
|
size_t getSerializationSize() const noexcept override;
|
void serialize(void* buffer) const noexcept override;
|
const char* getPluginType() const noexcept override
|
{
|
return "DETECT_TRT";
|
}
|
const char* getPluginVersion() const noexcept override
|
{
|
return "1.0";
|
}
|
void destroy() noexcept override
|
{
|
delete this;
|
}
|
void setPluginNamespace(const char* pluginNamespace) noexcept override
|
{
|
_s_plugin_namespace = pluginNamespace;
|
}
|
const char* getPluginNamespace() const noexcept override
|
{
|
return _s_plugin_namespace.c_str();
|
}
|
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept
|
{
|
return DataType::kFLOAT;
|
}
|
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const noexcept
|
{
|
return false;
|
}
|
bool canBroadcastInputAcrossBatch(int inputIndex) const noexcept
|
{
|
return false;
|
}
|
void attachToContext(
|
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)
|
{}
|
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) ;
|
void detachFromContext()
|
{}
|
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const noexcept
|
{
|
return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
|
}
|
IPluginV2* clone() const noexcept override;
|
private:
|
|
uint32_t _n_anchor;
|
uint32_t _n_classes;
|
uint32_t _n_grid_h;
|
uint32_t _n_grid_w;
|
//uint32_t _n_stride_h;
|
// uint32_t _n_stride_w;
|
uint64_t _n_output_size;
|
std::string _s_plugin_namespace;
|
}; //end detect
|
|
class DetectPluginCreator : public IPluginCreator
|
{
|
public:
|
DetectPluginCreator();
|
~DetectPluginCreator() override = default;
|
const char* getPluginName()const noexcept override;
|
const char* getPluginVersion() const noexcept override;
|
const PluginFieldCollection* getFieldNames() noexcept override;
|
IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override;
|
IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;
|
void setPluginNamespace(const char* libNamespace) noexcept override;
|
const char* getPluginNamespace() const noexcept override;
|
private:
|
std::string _s_name_space;
|
static PluginFieldCollection _fc;
|
static std::vector<PluginField> _vec_plugin_attributes;
|
};//end detect creator
|
|
}//end namespace nvinfer1
|
|
|
|
#endif
|
