151 lines
5.4 KiB
C++
151 lines
5.4 KiB
C++
#ifndef _YOLO_LAYER_H
|
|
#define _YOLO_LAYER_H
|
|
|
|
#include <cassert>
|
|
#include <vector>
|
|
#include <string>
|
|
#include <iostream>
|
|
#include "math_constants.h"
|
|
#include "NvInfer.h"
|
|
|
|
#define MAX_ANCHORS 6
|
|
|
|
#if NV_TENSORRT_MAJOR >= 8
|
|
#define NOEXCEPT noexcept
|
|
#else
|
|
#define NOEXCEPT
|
|
#endif
|
|
|
|
#define CHECK(status) \
|
|
do { \
|
|
auto ret = status; \
|
|
if (ret != 0) { \
|
|
std::cerr << "Cuda failure in file '" << __FILE__ \
|
|
<< "' line " << __LINE__ \
|
|
<< ": " << ret << std::endl; \
|
|
abort(); \
|
|
} \
|
|
} while (0)
|
|
|
|
namespace Yolo
|
|
{
|
|
static constexpr float IGNORE_THRESH = 0.01f;
|
|
|
|
struct alignas(float) Detection {
|
|
float bbox[4]; // x, y, w, h
|
|
float det_confidence;
|
|
float class_id;
|
|
float class_confidence;
|
|
};
|
|
}
|
|
|
|
namespace nvinfer1
|
|
{
|
|
class YoloLayerPlugin: public IPluginV2IOExt
|
|
{
|
|
public:
|
|
YoloLayerPlugin(int yolo_width, int yolo_height, int num_anchors, float* anchors, int num_classes, int input_width, int input_height, float scale_x_y, int new_coords);
|
|
YoloLayerPlugin(const void* data, size_t length);
|
|
|
|
~YoloLayerPlugin() override = default;
|
|
|
|
IPluginV2IOExt* clone() const NOEXCEPT override;
|
|
|
|
int initialize() NOEXCEPT override { return 0; }
|
|
|
|
void terminate() NOEXCEPT override;
|
|
|
|
void destroy() NOEXCEPT override { delete this; }
|
|
|
|
size_t getSerializationSize() const NOEXCEPT override;
|
|
|
|
void serialize(void* buffer) const NOEXCEPT override;
|
|
|
|
int getNbOutputs() const NOEXCEPT override { return 1; }
|
|
|
|
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override;
|
|
|
|
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0; }
|
|
|
|
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT override { return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; }
|
|
|
|
const char* getPluginType() const NOEXCEPT override { return "YoloLayer_TRT"; }
|
|
|
|
const char* getPluginVersion() const NOEXCEPT override { return "1"; }
|
|
|
|
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override { mPluginNamespace = pluginNamespace; }
|
|
|
|
const char* getPluginNamespace() const NOEXCEPT override { return mPluginNamespace; }
|
|
|
|
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const NOEXCEPT override { return DataType::kFLOAT; }
|
|
|
|
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override { return false; }
|
|
|
|
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override { return false; }
|
|
|
|
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override {}
|
|
|
|
//using IPluginV2IOExt::configurePlugin;
|
|
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) NOEXCEPT override {}
|
|
|
|
void detachFromContext() NOEXCEPT override {}
|
|
|
|
#if NV_TENSORRT_MAJOR >= 8
|
|
int32_t enqueue(int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override;
|
|
#else
|
|
int enqueue(int batchSize, const void* const * inputs, void** outputs, void* workspace, cudaStream_t stream) NOEXCEPT override;
|
|
#endif
|
|
|
|
private:
|
|
void forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int batchSize = 1);
|
|
|
|
int mThreadCount = 64;
|
|
int mYoloWidth, mYoloHeight, mNumAnchors;
|
|
float mAnchorsHost[MAX_ANCHORS * 2];
|
|
float *mAnchors; // allocated on GPU
|
|
int mNumClasses;
|
|
int mInputWidth, mInputHeight;
|
|
float mScaleXY;
|
|
int mNewCoords = 0;
|
|
|
|
const char* mPluginNamespace;
|
|
};
|
|
|
|
class YoloPluginCreator : public IPluginCreator
|
|
{
|
|
public:
|
|
YoloPluginCreator();
|
|
|
|
~YoloPluginCreator() override = default;
|
|
|
|
const char* getPluginName() const NOEXCEPT override;
|
|
|
|
const char* getPluginVersion() const NOEXCEPT override;
|
|
|
|
const PluginFieldCollection* getFieldNames() NOEXCEPT override;
|
|
|
|
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) NOEXCEPT override;
|
|
|
|
IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) NOEXCEPT override;
|
|
|
|
void setPluginNamespace(const char* libNamespace) NOEXCEPT override
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
const char* getPluginNamespace() const NOEXCEPT override
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
private:
|
|
static PluginFieldCollection mFC;
|
|
static std::vector<PluginField> mPluginAttributes;
|
|
std::string mNamespace;
|
|
};
|
|
|
|
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
|
|
};
|
|
|
|
#endif
|