thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

GPU Batched NMS

Browse Source
This commit is contained in:
Marcos Luciano
2022-06-19 03:25:50 -03:00
parent 2300e3b44b
commit f621c0f429
24 changed files with 835 additions and 654 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
nvdsinfer_custom_impl_Yolo/yoloPlugins.h
View File

@@ -36,14 +36,16 @@
#include "NvInferPlugin.h"
#define CUDA_CHECK(status) \
{ \
if (status != 0) \
{ \
std::cout << "CUDA failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \
<< " at line " << __LINE__ << std::endl; \
abort(); \
} \
#include "yolo.h"
#define CUDA_CHECK(status) \
{ \
if (status != 0) \
{ \
std::cout << "CUDA failure: " << cudaGetErrorString(status) << " in file " << __FILE__ << " at line " \
<< __LINE__ << std::endl; \
abort(); \
} \
}
namespace
@@ -56,15 +58,17 @@ class YoloLayer : public nvinfer1::IPluginV2
{
public:
YoloLayer (const void* data, size_t length);
YoloLayer (
const uint& numBBoxes, const uint& numClasses, const uint& netWidth, const uint& netHeight,
const uint& gridSizeX, const uint& gridSizeY, const uint& modelType, const uint& newCoords,
const float& scaleXY, const std::vector<float> anchors,
const std::vector<int> mask);
~YoloLayer ();
const uint& netWidth, const uint& netHeight, const uint& numClasses, const uint& newCoords,
const std::vector<TensorInfo>& yoloTensors, const uint64_t& outputSize, const uint& modelType, const uint& topK,
const float& scoreThreshold);
const char* getPluginType () const noexcept override { return YOLOLAYER_PLUGIN_NAME; }
const char* getPluginVersion () const noexcept override { return YOLOLAYER_PLUGIN_VERSION; }
int getNbOutputs () const noexcept override { return 1; }
int getNbOutputs () const noexcept override { return 2; }
nvinfer1::Dims getOutputDimensions (
int index, const nvinfer1::Dims* inputs,
@@ -74,53 +78,59 @@ public:
nvinfer1::DataType type, nvinfer1::PluginFormat format) const noexcept override;
void configureWithFormat (
const nvinfer1::Dims* inputDims, int nbInputs,
const nvinfer1::Dims* outputDims, int nbOutputs,
const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims, int nbOutputs,
nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) noexcept override;
int initialize () noexcept override { return 0; }
void terminate () noexcept override {}
size_t getWorkspaceSize (int maxBatchSize) const noexcept override { return 0; }
size_t getWorkspaceSize (int maxBatchSize) const noexcept override {
return maxBatchSize * sizeof(int);
}
int32_t enqueue (
int32_t batchSize, void const* const* inputs, void* const* outputs,
void* workspace, cudaStream_t stream) noexcept override;
int batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream)
noexcept override;
size_t getSerializationSize() const noexcept override;
void serialize (void* buffer) const noexcept override;
void destroy () noexcept override { delete this; }
nvinfer1::IPluginV2* clone() const noexcept override;
void setPluginNamespace (const char* pluginNamespace) noexcept override {
m_Namespace = pluginNamespace;
}
virtual const char* getPluginNamespace () const noexcept override {
return m_Namespace.c_str();
}
private:
std::string m_Namespace {""};
uint m_NumBBoxes {0};
uint m_NumClasses {0};
uint m_NetWidth {0};
uint m_NetHeight {0};
uint m_GridSizeX {0};
uint m_GridSizeY {0};
uint m_Type {0};
uint m_NumClasses {0};
uint m_NewCoords {0};
float m_ScaleXY {0};
std::vector<float> m_Anchors;
std::vector<int> m_Mask;
std::vector<TensorInfo> m_YoloTensors;
uint64_t m_OutputSize {0};
void* p_Anchors;
void* p_Mask;
uint m_Type {0};
uint m_TopK {0};
float m_ScoreThreshold {0};
};
class YoloLayerPluginCreator : public nvinfer1::IPluginCreator
{
public:
YoloLayerPluginCreator () {}
~YoloLayerPluginCreator () {}
const char* getPluginName () const noexcept override { return YOLOLAYER_PLUGIN_NAME; }
const char* getPluginVersion () const noexcept override { return YOLOLAYER_PLUGIN_VERSION; }
const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override {
@@ -153,7 +163,6 @@ private:
std::string m_Namespace {""};
};
extern uint kNUM_BBOXES;
extern uint kNUM_CLASSES;
#endif // __YOLO_PLUGINS__