diff --git a/non_square/nvdsinfer_custom_impl_Yolo/Makefile b/non_square/nvdsinfer_custom_impl_Yolo/Makefile new file mode 100644 index 0000000..1780210 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/Makefile @@ -0,0 +1,71 @@ +################################################################################ +# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +# Edited by Marcos Luciano +# https://www.github.com/marcoslucianops +################################################################################ + +CUDA_VER?= +ifeq ($(CUDA_VER),) + $(error "CUDA_VER is not set") +endif +CC:= g++ +NVCC:=/usr/local/cuda-$(CUDA_VER)/bin/nvcc + +CFLAGS:= -Wall -std=c++11 -shared -fPIC -Wno-error=deprecated-declarations +CFLAGS+= -I../../includes -I/usr/local/cuda-$(CUDA_VER)/include + +LIBS:= -lnvinfer_plugin -lnvinfer -lnvparsers -L/usr/local/cuda-$(CUDA_VER)/lib64 -lcudart -lcublas -lstdc++fs +LFLAGS:= -shared -Wl,--start-group $(LIBS) -Wl,--end-group + +INCS:= $(wildcard *.h) +SRCFILES:= nvdsinfer_yolo_engine.cpp \ + nvdsparsebbox_Yolo.cpp \ + yoloPlugins.cpp \ + layers/convolutional_layer.cpp \ + layers/dropout_layer.cpp \ + layers/shortcut_layer.cpp \ + layers/route_layer.cpp \ + layers/upsample_layer.cpp \ + layers/maxpool_layer.cpp \ + layers/activation_layer.cpp \ + utils.cpp \ + yolo.cpp \ + yoloForward.cu +TARGET_LIB:= libnvdsinfer_custom_impl_Yolo.so + +TARGET_OBJS:= $(SRCFILES:.cpp=.o) +TARGET_OBJS:= $(TARGET_OBJS:.cu=.o) + +all: $(TARGET_LIB) + +%.o: %.cpp $(INCS) Makefile + $(CC) -c -o $@ $(CFLAGS) $< + +%.o: %.cu $(INCS) Makefile + $(NVCC) -c -o $@ --compiler-options '-fPIC' $< + +$(TARGET_LIB) : $(TARGET_OBJS) + $(CC) -o $@ $(TARGET_OBJS) $(LFLAGS) + +clean: + rm -rf $(TARGET_LIB) + rm -rf $(TARGET_OBJS) diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp new file mode 100644 index 0000000..d730fd2 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp @@ -0,0 +1,82 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "activation_layer.h" + +nvinfer1::ILayer* activationLayer( + int layerIdx, + std::string activation, + nvinfer1::ILayer* output, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network) +{ + if (activation == "relu") + { + nvinfer1::IActivationLayer* relu = network->addActivation( + *input, nvinfer1::ActivationType::kRELU); + assert(relu != nullptr); + std::string reluLayerName = "relu_" + std::to_string(layerIdx); + relu->setName(reluLayerName.c_str()); + output = relu; + } + else if (activation == "sigmoid" || activation == "logistic") + { + nvinfer1::IActivationLayer* sigmoid = network->addActivation( + *input, nvinfer1::ActivationType::kSIGMOID); + assert(sigmoid != nullptr); + std::string sigmoidLayerName = "sigmoid_" + std::to_string(layerIdx); + sigmoid->setName(sigmoidLayerName.c_str()); + output = sigmoid; + } + else if (activation == "tanh") + { + nvinfer1::IActivationLayer* tanh = network->addActivation( + *input, nvinfer1::ActivationType::kTANH); + assert(tanh != nullptr); + std::string tanhLayerName = "tanh_" + std::to_string(layerIdx); + tanh->setName(tanhLayerName.c_str()); + output = tanh; + } + else if (activation == "leaky") + { + nvinfer1::IActivationLayer* leaky = network->addActivation( + *input, nvinfer1::ActivationType::kLEAKY_RELU); + leaky->setAlpha(0.1); + assert(leaky != nullptr); + std::string leakyLayerName = "leaky_" + std::to_string(layerIdx); + leaky->setName(leakyLayerName.c_str()); + output = leaky; + } + else if (activation == "softplus") + { + nvinfer1::IActivationLayer* softplus = network->addActivation( + *input, nvinfer1::ActivationType::kSOFTPLUS); + assert(softplus != nullptr); + std::string softplusLayerName = "softplus_" + std::to_string(layerIdx); + softplus->setName(softplusLayerName.c_str()); + output = softplus; + } + else if (activation == "mish") + { + nvinfer1::IActivationLayer* softplus = network->addActivation( + *input, nvinfer1::ActivationType::kSOFTPLUS); + assert(softplus != nullptr); + std::string softplusLayerName = "softplus_" + std::to_string(layerIdx); + softplus->setName(softplusLayerName.c_str()); + nvinfer1::IActivationLayer* tanh = network->addActivation( + *softplus->getOutput(0), nvinfer1::ActivationType::kTANH); + assert(tanh != nullptr); + std::string tanhLayerName = "tanh_" + std::to_string(layerIdx); + tanh->setName(tanhLayerName.c_str()); + nvinfer1::IElementWiseLayer* mish = network->addElementWise( + *tanh->getOutput(0), *input, + nvinfer1::ElementWiseOperation::kPROD); + assert(mish != nullptr); + std::string mishLayerName = "mish_" + std::to_string(layerIdx); + mish->setName(mishLayerName.c_str()); + output = mish; + } + return output; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/activation_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/activation_layer.h new file mode 100644 index 0000000..e6081e6 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/activation_layer.h @@ -0,0 +1,23 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __ACTIVATION_LAYER_H__ +#define __ACTIVATION_LAYER_H__ + +#include +#include + +#include "NvInfer.h" + +#include "activation_layer.h" + +nvinfer1::ILayer* activationLayer( + int layerIdx, + std::string activation, + nvinfer1::ILayer* output, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp new file mode 100644 index 0000000..abb0d32 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp @@ -0,0 +1,168 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include +#include "convolutional_layer.h" + +nvinfer1::ILayer* convolutionalLayer( + int layerIdx, + std::map& block, + std::vector& weights, + std::vector& trtWeights, + int& weightPtr, + int& inputChannels, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network) +{ + assert(block.at("type") == "convolutional"); + assert(block.find("filters") != block.end()); + assert(block.find("pad") != block.end()); + assert(block.find("size") != block.end()); + assert(block.find("stride") != block.end()); + + int filters = std::stoi(block.at("filters")); + int padding = std::stoi(block.at("pad")); + int kernelSize = std::stoi(block.at("size")); + int stride = std::stoi(block.at("stride")); + std::string activation = block.at("activation"); + int bias = filters; + + bool batchNormalize = false; + if (block.find("batch_normalize") != block.end()) + { + bias = 0; + batchNormalize = (block.at("batch_normalize") == "1"); + } + + int groups = 1; + if (block.find("groups") != block.end()) + { + groups = std::stoi(block.at("groups")); + } + + int pad; + if (padding) + pad = (kernelSize - 1) / 2; + else + pad = 0; + + int size = filters * inputChannels * kernelSize * kernelSize / groups; + std::vector bnBiases; + std::vector bnWeights; + std::vector bnRunningMean; + std::vector bnRunningVar; + nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size}; + nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, bias}; + + if (batchNormalize == false) + { + float* val = new float[filters]; + for (int i = 0; i < filters; ++i) + { + val[i] = weights[weightPtr]; + weightPtr++; + } + convBias.values = val; + trtWeights.push_back(convBias); + val = new float[size]; + for (int i = 0; i < size; ++i) + { + val[i] = weights[weightPtr]; + weightPtr++; + } + convWt.values = val; + trtWeights.push_back(convWt); + } + else + { + for (int i = 0; i < filters; ++i) + { + bnBiases.push_back(weights[weightPtr]); + weightPtr++; + } + + for (int i = 0; i < filters; ++i) + { + bnWeights.push_back(weights[weightPtr]); + weightPtr++; + } + for (int i = 0; i < filters; ++i) + { + bnRunningMean.push_back(weights[weightPtr]); + weightPtr++; + } + for (int i = 0; i < filters; ++i) + { + bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5)); + weightPtr++; + } + float* val = new float[size]; + for (int i = 0; i < size; ++i) + { + val[i] = weights[weightPtr]; + weightPtr++; + } + convWt.values = val; + trtWeights.push_back(convWt); + trtWeights.push_back(convBias); + } + + nvinfer1::IConvolutionLayer* conv = network->addConvolution( + *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias); + assert(conv != nullptr); + std::string convLayerName = "conv_" + std::to_string(layerIdx); + conv->setName(convLayerName.c_str()); + conv->setStride(nvinfer1::DimsHW{stride, stride}); + conv->setPadding(nvinfer1::DimsHW{pad, pad}); + + if (block.find("groups") != block.end()) + { + conv->setNbGroups(groups); + } + + nvinfer1::ILayer* output = conv; + + if (batchNormalize == true) + { + size = filters; + nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size}; + nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size}; + nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size}; + float* shiftWt = new float[size]; + for (int i = 0; i < size; ++i) + { + shiftWt[i] + = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i)); + } + shift.values = shiftWt; + float* scaleWt = new float[size]; + for (int i = 0; i < size; ++i) + { + scaleWt[i] = bnWeights.at(i) / bnRunningVar[i]; + } + scale.values = scaleWt; + float* powerWt = new float[size]; + for (int i = 0; i < size; ++i) + { + powerWt[i] = 1.0; + } + power.values = powerWt; + trtWeights.push_back(shift); + trtWeights.push_back(scale); + trtWeights.push_back(power); + + nvinfer1::IScaleLayer* bn = network->addScale( + *output->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shift, scale, power); + assert(bn != nullptr); + std::string bnLayerName = "batch_norm_" + std::to_string(layerIdx); + bn->setName(bnLayerName.c_str()); + output = bn; + } + + output = activationLayer(layerIdx, activation, output, output->getOutput(0), network); + assert(output != nullptr); + + return output; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.h new file mode 100644 index 0000000..b114493 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.h @@ -0,0 +1,26 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __CONVOLUTIONAL_LAYER_H__ +#define __CONVOLUTIONAL_LAYER_H__ + +#include +#include + +#include "NvInfer.h" + +#include "activation_layer.h" + +nvinfer1::ILayer* convolutionalLayer( + int layerIdx, + std::map& block, + std::vector& weights, + std::vector& trtWeights, + int& weightPtr, + int& inputChannels, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.cpp new file mode 100644 index 0000000..112af16 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.cpp @@ -0,0 +1,15 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "dropout_layer.h" + +nvinfer1::ILayer* dropoutLayer( + float probability, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network) +{ + nvinfer1::ILayer* output; + return output; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.h new file mode 100644 index 0000000..3e8a1d4 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.h @@ -0,0 +1,16 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __DROPOUT_LAYER_H__ +#define __DROPOUT_LAYER_H__ + +#include "NvInfer.h" + +nvinfer1::ILayer* dropoutLayer( + float probability, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp new file mode 100644 index 0000000..06948dc --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp @@ -0,0 +1,30 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "maxpool_layer.h" + +nvinfer1::ILayer* maxpoolLayer( + int layerIdx, + std::map& block, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network) +{ + assert(block.at("type") == "maxpool"); + assert(block.find("size") != block.end()); + assert(block.find("stride") != block.end()); + + int size = std::stoi(block.at("size")); + int stride = std::stoi(block.at("stride")); + + nvinfer1::IPoolingLayer* pool + = network->addPooling(*input, nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{size, size}); + assert(pool); + std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx); + pool->setStride(nvinfer1::DimsHW{stride, stride}); + pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER); + pool->setName(maxpoolLayerName.c_str()); + + return pool; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.h new file mode 100644 index 0000000..a0d139e --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.h @@ -0,0 +1,20 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __MAXPOOL_LAYER_H__ +#define __MAXPOOL_LAYER_H__ + +#include +#include + +#include "NvInfer.h" + +nvinfer1::ILayer* maxpoolLayer( + int layerIdx, + std::map& block, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp new file mode 100644 index 0000000..d4bf5dc --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp @@ -0,0 +1,63 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "route_layer.h" + +nvinfer1::ILayer* routeLayer( + int layerIdx, + std::map& block, + std::vector tensorOutputs, + nvinfer1::INetworkDefinition* network) +{ + std::string strLayers = block.at("layers"); + std::vector idxLayers; + size_t lastPos = 0, pos = 0; + while ((pos = strLayers.find(',', lastPos)) != std::string::npos) { + int vL = std::stoi(trim(strLayers.substr(lastPos, pos - lastPos))); + idxLayers.push_back (vL); + lastPos = pos + 1; + } + if (lastPos < strLayers.length()) { + std::string lastV = trim(strLayers.substr(lastPos)); + if (!lastV.empty()) { + idxLayers.push_back (std::stoi(lastV)); + } + } + assert (!idxLayers.empty()); + std::vector concatInputs; + for (int idxLayer : idxLayers) { + if (idxLayer < 0) { + idxLayer = tensorOutputs.size() + idxLayer; + } + assert (idxLayer >= 0 && idxLayer < (int)tensorOutputs.size()); + concatInputs.push_back (tensorOutputs[idxLayer]); + } + + nvinfer1::IConcatenationLayer* concat = + network->addConcatenation(concatInputs.data(), concatInputs.size()); + assert(concat != nullptr); + std::string concatLayerName = "route_" + std::to_string(layerIdx - 1); + concat->setName(concatLayerName.c_str()); + concat->setAxis(0); + + nvinfer1::ILayer* output = concat; + + if (block.find("groups") != block.end()) { + nvinfer1::Dims prevTensorDims = output->getOutput(0)->getDimensions(); + int groups = stoi(block.at("groups")); + int group_id = stoi(block.at("group_id")); + int startSlice = (prevTensorDims.d[0] / groups) * group_id; + int channelSlice = (prevTensorDims.d[0] / groups); + nvinfer1::ISliceLayer* sl = network->addSlice( + *output->getOutput(0), + nvinfer1::Dims3{startSlice, 0, 0}, + nvinfer1::Dims3{channelSlice, prevTensorDims.d[1], prevTensorDims.d[2]}, + nvinfer1::Dims3{1, 1, 1}); + assert(sl != nullptr); + output = sl; + } + + return output; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/route_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/route_layer.h new file mode 100644 index 0000000..63c6bd6 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/route_layer.h @@ -0,0 +1,18 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __ROUTE_LAYER_H__ +#define __ROUTE_LAYER_H__ + +#include "NvInfer.h" +#include "../utils.h" + +nvinfer1::ILayer* routeLayer( + int layerIdx, + std::map& block, + std::vector tensorOutputs, + nvinfer1::INetworkDefinition* network); + +#endif \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp new file mode 100644 index 0000000..87ba3f8 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp @@ -0,0 +1,45 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "shortcut_layer.h" + +nvinfer1::ILayer* shortcutLayer( + int layerIdx, + std::string activation, + std::string inputVol, + std::string shortcutVol, + nvinfer1::ITensor* input, + nvinfer1::ITensor* shortcutTensor, + nvinfer1::INetworkDefinition* network) +{ + nvinfer1::ILayer* output; + nvinfer1::ITensor* outputTensor; + + if (inputVol != shortcutVol) + { + nvinfer1::ISliceLayer* sl = network->addSlice( + *shortcutTensor, + nvinfer1::Dims3{0, 0, 0}, + input->getDimensions(), + nvinfer1::Dims3{1, 1, 1}); + assert(sl != nullptr); + outputTensor = sl->getOutput(0); + assert(outputTensor != nullptr); + } else + { + outputTensor = shortcutTensor; + assert(outputTensor != nullptr); + } + + nvinfer1::IElementWiseLayer* ew = network->addElementWise( + *input, *outputTensor, + nvinfer1::ElementWiseOperation::kSUM); + assert(ew != nullptr); + + output = activationLayer(layerIdx, activation, ew, ew->getOutput(0), network); + assert(output != nullptr); + + return output; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.h new file mode 100644 index 0000000..c564d3d --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.h @@ -0,0 +1,22 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __SHORTCUT_LAYER_H__ +#define __SHORTCUT_LAYER_H__ + +#include "NvInfer.h" + +#include "activation_layer.h" + +nvinfer1::ILayer* shortcutLayer( + int layerIdx, + std::string activation, + std::string inputVol, + std::string shortcutVol, + nvinfer1::ITensor* input, + nvinfer1::ITensor* shortcutTensor, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp b/non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp new file mode 100644 index 0000000..7dc0a26 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp @@ -0,0 +1,83 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "upsample_layer.h" + +nvinfer1::ILayer* upsampleLayer( + int layerIdx, + std::map& block, + std::vector& weights, + std::vector& trtWeights, + int& inputChannels, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network) +{ + assert(block.at("type") == "upsample"); + nvinfer1::Dims inpDims = input->getDimensions(); + assert(inpDims.nbDims == 3); + //assert(inpDims.d[1] == inpDims.d[2]); + int h = inpDims.d[1]; + int w = inpDims.d[2]; + int stride = std::stoi(block.at("stride")); + + nvinfer1::Dims preDims{3, + {1, stride * h, h}, + {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, + nvinfer1::DimensionType::kSPATIAL}}; + int size = stride * h * w; + nvinfer1::Weights preMul{nvinfer1::DataType::kFLOAT, nullptr, size}; + float* preWt = new float[size]; + + for (int i = 0, idx = 0; i < h; ++i) + { + for (int j = 0; j < h * stride; ++j, ++idx) + { + preWt[idx] = (i == j) ? 1.0 : 0.0; + } + } + preMul.values = preWt; + trtWeights.push_back(preMul); + nvinfer1::IConstantLayer* preM = network->addConstant(preDims, preMul); + assert(preM != nullptr); + std::string preLayerName = "preMul_" + std::to_string(layerIdx); + preM->setName(preLayerName.c_str()); + + nvinfer1::Dims postDims{3, + {1, w, stride * w}, + {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, + nvinfer1::DimensionType::kSPATIAL}}; + size = stride * w * w; + nvinfer1::Weights postMul{nvinfer1::DataType::kFLOAT, nullptr, size}; + float* postWt = new float[size]; + + for (int i = 0, idx = 0; i < w; ++i) + { + for (int j = 0; j < stride * w; ++j, ++idx) + { + postWt[idx] = (j / stride == i) ? 1.0 : 0.0; + } + } + postMul.values = postWt; + trtWeights.push_back(postMul); + nvinfer1::IConstantLayer* post_m = network->addConstant(postDims, postMul); + assert(post_m != nullptr); + std::string postLayerName = "postMul_" + std::to_string(layerIdx); + post_m->setName(postLayerName.c_str()); + + nvinfer1::IMatrixMultiplyLayer* mm1 + = network->addMatrixMultiply(*preM->getOutput(0), nvinfer1::MatrixOperation::kNONE, *input, + nvinfer1::MatrixOperation::kNONE); + assert(mm1 != nullptr); + std::string mm1LayerName = "mm1_" + std::to_string(layerIdx); + mm1->setName(mm1LayerName.c_str()); + nvinfer1::IMatrixMultiplyLayer* mm2 + = network->addMatrixMultiply(*mm1->getOutput(0), nvinfer1::MatrixOperation::kNONE, + *post_m->getOutput(0), nvinfer1::MatrixOperation::kNONE); + assert(mm2 != nullptr); + std::string mm2LayerName = "mm2_" + std::to_string(layerIdx); + mm2->setName(mm2LayerName.c_str()); + + return mm2; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.h b/non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.h new file mode 100644 index 0000000..6b0224a --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.h @@ -0,0 +1,24 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __UPSAMPLE_LAYER_H__ +#define __UPSAMPLE_LAYER_H__ + +#include +#include +#include + +#include "NvInfer.h" + +nvinfer1::ILayer* upsampleLayer( + int layerIdx, + std::map& block, + std::vector& weights, + std::vector& trtWeights, + int& inputChannels, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp b/non_square/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp new file mode 100644 index 0000000..29306cd --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp @@ -0,0 +1,107 @@ +/* + * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "nvdsinfer_custom_impl.h" +#include "nvdsinfer_context.h" +#include "yoloPlugins.h" +#include "yolo.h" + +#include + +#define USE_CUDA_ENGINE_GET_API 1 + +static bool getYoloNetworkInfo (NetworkInfo &networkInfo, const NvDsInferContextInitParams* initParams) +{ + std::string yoloCfg = initParams->customNetworkConfigFilePath; + std::string yoloType; + + std::transform (yoloCfg.begin(), yoloCfg.end(), yoloCfg.begin(), [] (uint8_t c) { + return std::tolower (c);}); + + yoloType = yoloCfg.substr(0, yoloCfg.find(".cfg")); + + networkInfo.networkType = yoloType; + networkInfo.configFilePath = initParams->customNetworkConfigFilePath; + networkInfo.wtsFilePath = initParams->modelFilePath; + networkInfo.deviceType = (initParams->useDLA ? "kDLA" : "kGPU"); + networkInfo.inputBlobName = "data"; + + if (networkInfo.configFilePath.empty() || + networkInfo.wtsFilePath.empty()) { + std::cerr << "YOLO config file or weights file is not specified" + << std::endl; + return false; + } + + if (!fileExists(networkInfo.configFilePath) || + !fileExists(networkInfo.wtsFilePath)) { + std::cerr << "YOLO config file or weights file is not exist" + << std::endl; + return false; + } + + return true; +} + +#if !USE_CUDA_ENGINE_GET_API +IModelParser* NvDsInferCreateModelParser( + const NvDsInferContextInitParams* initParams) { + NetworkInfo networkInfo; + if (!getYoloNetworkInfo(networkInfo, initParams)) { + return nullptr; + } + + return new Yolo(networkInfo); +} +#else +extern "C" +bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, + const NvDsInferContextInitParams * const initParams, + nvinfer1::DataType dataType, + nvinfer1::ICudaEngine *& cudaEngine); + +extern "C" +bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, + const NvDsInferContextInitParams * const initParams, + nvinfer1::DataType dataType, + nvinfer1::ICudaEngine *& cudaEngine) +{ + NetworkInfo networkInfo; + if (!getYoloNetworkInfo(networkInfo, initParams)) { + return false; + } + + Yolo yolo(networkInfo); + cudaEngine = yolo.createEngine (builder); + if (cudaEngine == nullptr) + { + std::cerr << "Failed to build CUDA engine on " + << networkInfo.configFilePath << std::endl; + return false; + } + + return true; +} +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp b/non_square/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp new file mode 100644 index 0000000..5c9801b --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp @@ -0,0 +1,380 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include +#include +#include +#include "nvdsinfer_custom_impl.h" +#include "utils.h" + +#include "yoloPlugins.h" + +extern "C" bool NvDsInferParseYolo( + std::vector const& outputLayersInfo, + NvDsInferNetworkInfo const& networkInfo, + NvDsInferParseDetectionParams const& detectionParams, + std::vector& objectList); + +static std::vector +nonMaximumSuppression(const float nmsThresh, std::vector binfo) +{ + auto overlap1D = [](float x1min, float x1max, float x2min, float x2max) -> float { + if (x1min > x2min) + { + std::swap(x1min, x2min); + std::swap(x1max, x2max); + } + return x1max < x2min ? 0 : std::min(x1max, x2max) - x2min; + }; + auto computeIoU + = [&overlap1D](NvDsInferParseObjectInfo& bbox1, NvDsInferParseObjectInfo& bbox2) -> float { + float overlapX + = overlap1D(bbox1.left, bbox1.left + bbox1.width, bbox2.left, bbox2.left + bbox2.width); + float overlapY + = overlap1D(bbox1.top, bbox1.top + bbox1.height, bbox2.top, bbox2.top + bbox2.height); + float area1 = (bbox1.width) * (bbox1.height); + float area2 = (bbox2.width) * (bbox2.height); + float overlap2D = overlapX * overlapY; + float u = area1 + area2 - overlap2D; + return u == 0 ? 0 : overlap2D / u; + }; + + std::stable_sort(binfo.begin(), binfo.end(), + [](const NvDsInferParseObjectInfo& b1, const NvDsInferParseObjectInfo& b2) { + return b1.detectionConfidence > b2.detectionConfidence; + }); + std::vector out; + for (auto i : binfo) + { + bool keep = true; + for (auto j : out) + { + if (keep) + { + float overlap = computeIoU(i, j); + keep = overlap <= nmsThresh; + } + else + break; + } + if (keep) out.push_back(i); + } + return out; +} + +static std::vector +nmsAllClasses(const float nmsThresh, + std::vector& binfo, + const uint numClasses) +{ + std::vector result; + std::vector> splitBoxes(numClasses); + for (auto& box : binfo) + { + splitBoxes.at(box.classId).push_back(box); + } + + for (auto& boxes : splitBoxes) + { + boxes = nonMaximumSuppression(nmsThresh, boxes); + result.insert(result.end(), boxes.begin(), boxes.end()); + } + return result; +} + +static NvDsInferParseObjectInfo convertBBox(const float& bx, const float& by, const float& bw, + const float& bh, const int& stride, const uint& netW, + const uint& netH) +{ + NvDsInferParseObjectInfo b; + float xCenter = bx * stride; + float yCenter = by * stride; + float x0 = xCenter - bw / 2; + float y0 = yCenter - bh / 2; + float x1 = x0 + bw; + float y1 = y0 + bh; + + x0 = clamp(x0, 0, netW); + y0 = clamp(y0, 0, netH); + x1 = clamp(x1, 0, netW); + y1 = clamp(y1, 0, netH); + + b.left = x0; + b.width = clamp(x1 - x0, 0, netW); + b.top = y0; + b.height = clamp(y1 - y0, 0, netH); + + return b; +} + +static void addBBoxProposal(const float bx, const float by, const float bw, const float bh, + const uint stride, const uint& netW, const uint& netH, const int maxIndex, + const float maxProb, std::vector& binfo) +{ + NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, stride, netW, netH); + if (bbi.width < 1 || bbi.height < 1) return; + + bbi.detectionConfidence = maxProb; + bbi.classId = maxIndex; + binfo.push_back(bbi); +} + +static std::vector +decodeYoloTensor( + const float* detections, const std::vector &mask, const std::vector &anchors, + const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, + const uint numOutputClasses, const uint& netW, + const uint& netH, + const float confThresh) +{ + std::vector binfo; + for (uint y = 0; y < gridSizeH; ++y) { + for (uint x = 0; x < gridSizeW; ++x) { + for (uint b = 0; b < numBBoxes; ++b) + { + const float pw = anchors[mask[b] * 2]; + const float ph = anchors[mask[b] * 2 + 1]; + + const int numGridCells = gridSizeH * gridSizeW; + const int bbindex = y * gridSizeW + x; + const float bx + = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; + const float by + = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; + const float bw + = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]; + const float bh + = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]; + + const float objectness + = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; + + float maxProb = 0.0f; + int maxIndex = -1; + + for (uint i = 0; i < numOutputClasses; ++i) + { + float prob + = (detections[bbindex + + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); + + if (prob > maxProb) + { + maxProb = prob; + maxIndex = i; + } + } + maxProb = objectness * maxProb; + + if (maxProb > confThresh) + { + addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); + } + } + } + } + return binfo; +} + +static std::vector +decodeYoloV2Tensor( + const float* detections, const std::vector &anchors, + const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, + const uint numOutputClasses, const uint& netW, + const uint& netH) +{ + std::vector binfo; + for (uint y = 0; y < gridSizeH; ++y) { + for (uint x = 0; x < gridSizeW; ++x) { + for (uint b = 0; b < numBBoxes; ++b) + { + const float pw = anchors[b * 2]; + const float ph = anchors[b * 2 + 1]; + + const int numGridCells = gridSizeH * gridSizeW; + const int bbindex = y * gridSizeW + x; + const float bx + = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; + const float by + = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; + const float bw + = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]; + const float bh + = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]; + + const float objectness + = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; + + float maxProb = 0.0f; + int maxIndex = -1; + + for (uint i = 0; i < numOutputClasses; ++i) + { + float prob + = (detections[bbindex + + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); + + if (prob > maxProb) + { + maxProb = prob; + maxIndex = i; + } + } + maxProb = objectness * maxProb; + + addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); + } + } + } + return binfo; +} + +static inline std::vector +SortLayers(const std::vector & outputLayersInfo) +{ + std::vector outLayers; + for (auto const &layer : outputLayersInfo) { + outLayers.push_back (&layer); + } + std::sort(outLayers.begin(), outLayers.end(), + [](const NvDsInferLayerInfo* a, const NvDsInferLayerInfo* b) { + return a->inferDims.d[1] < b->inferDims.d[1]; + }); + return outLayers; +} + +static bool NvDsInferParseYolo( + std::vector const& outputLayersInfo, + NvDsInferNetworkInfo const& networkInfo, + NvDsInferParseDetectionParams const& detectionParams, + std::vector& objectList, + const std::vector &anchors, + const std::vector> &masks, + const uint &num_classes, + const float &beta_nms) +{ + const float kCONF_THRESH = detectionParams.perClassThreshold[0]; + + const std::vector sortedLayers = + SortLayers (outputLayersInfo); + + if (sortedLayers.size() != masks.size()) { + std::cerr << "ERROR: YOLO output layer.size: " << sortedLayers.size() + << " does not match mask.size: " << masks.size() << std::endl; + return false; + } + + if (num_classes != detectionParams.numClassesConfigured) + { + std::cerr << "WARNING: Num classes mismatch. Configured: " + << detectionParams.numClassesConfigured + << ", detected by network: " << num_classes << std::endl; + } + + std::vector objects; + + for (uint idx = 0; idx < masks.size(); ++idx) { + const NvDsInferLayerInfo &layer = *sortedLayers[idx]; // 255 x Grid x Grid + + assert(layer.inferDims.numDims == 3); + const uint gridSizeH = layer.inferDims.d[1]; + const uint gridSizeW = layer.inferDims.d[2]; + const uint stride = DIVUP(networkInfo.width, gridSizeW); + //assert(stride == DIVUP(networkInfo.height, gridSizeH)); + + std::vector outObjs = + decodeYoloTensor((const float*)(layer.buffer), masks[idx], anchors, gridSizeW, gridSizeH, stride, masks[idx].size(), + num_classes, networkInfo.width, networkInfo.height, kCONF_THRESH); + objects.insert(objects.end(), outObjs.begin(), outObjs.end()); + } + + + objectList.clear(); + objectList = nmsAllClasses(beta_nms, objects, num_classes); + + return true; +} + +static bool NvDsInferParseYoloV2( + std::vector const& outputLayersInfo, + NvDsInferNetworkInfo const& networkInfo, + NvDsInferParseDetectionParams const& detectionParams, + std::vector& objectList, + std::vector &anchors, + const uint &num_classes) +{ + if (outputLayersInfo.empty()) { + std::cerr << "Could not find output layer in bbox parsing" << std::endl;; + return false; + } + const uint kNUM_BBOXES = anchors.size() / 2; + + const NvDsInferLayerInfo &layer = outputLayersInfo[0]; + + if (num_classes != detectionParams.numClassesConfigured) + { + std::cerr << "WARNING: Num classes mismatch. Configured: " + << detectionParams.numClassesConfigured + << ", detected by network: " << num_classes << std::endl; + } + + assert(layer.inferDims.numDims == 3); + const uint gridSizeH = layer.inferDims.d[1]; + const uint gridSizeW = layer.inferDims.d[2]; + const uint stride = DIVUP(networkInfo.width, gridSizeW); + //assert(stride == DIVUP(networkInfo.height, gridSizeH)); + for (auto& anchor : anchors) { + anchor *= stride; + } + std::vector objects = + decodeYoloV2Tensor((const float*)(layer.buffer), anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES, + num_classes, networkInfo.width, networkInfo.height); + + objectList = objects; + + return true; +} + +extern "C" bool NvDsInferParseYolo( + std::vector const& outputLayersInfo, + NvDsInferNetworkInfo const& networkInfo, + NvDsInferParseDetectionParams const& detectionParams, + std::vector& objectList) +{ + + int num_classes = kNUM_CLASSES; + float beta_nms = kBETA_NMS; + std::vector anchors = kANCHORS; + std::vector> mask = kMASK; + + if (mask.size() > 0) { + return NvDsInferParseYolo (outputLayersInfo, networkInfo, detectionParams, objectList, anchors, mask, num_classes, beta_nms); + } + else { + return NvDsInferParseYoloV2 (outputLayersInfo, networkInfo, detectionParams, objectList, anchors, num_classes); + } +} + +CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseYolo); \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/utils.cpp b/non_square/nvdsinfer_custom_impl_Yolo/utils.cpp new file mode 100644 index 0000000..5afe57e --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/utils.cpp @@ -0,0 +1,150 @@ +/* + * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "utils.h" + +#include +#include +#include +#include + +static void leftTrim(std::string& s) +{ + s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); })); +} + +static void rightTrim(std::string& s) +{ + s.erase(std::find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end()); +} + +std::string trim(std::string s) +{ + leftTrim(s); + rightTrim(s); + return s; +} + +float clamp(const float val, const float minVal, const float maxVal) +{ + assert(minVal <= maxVal); + return std::min(maxVal, std::max(minVal, val)); +} + +bool fileExists(const std::string fileName, bool verbose) +{ + if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName))) + { + if (verbose) std::cout << "File does not exist: " << fileName << std::endl; + return false; + } + return true; +} + +std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType) +{ + assert(fileExists(weightsFilePath)); + std::cout << "\nLoading pre-trained weights" << std::endl; + std::ifstream file(weightsFilePath, std::ios_base::binary); + assert(file.good()); + std::string line; + + if (networkType.find("yolov2") != std::string::npos && networkType.find("yolov2-tiny") == std::string::npos) + { + // Remove 4 int32 bytes of data from the stream belonging to the header + file.ignore(4 * 4); + } + else + { + // Remove 5 int32 bytes of data from the stream belonging to the header + file.ignore(4 * 5); + } + + std::vector weights; + char floatWeight[4]; + while (!file.eof()) + { + file.read(floatWeight, 4); + assert(file.gcount() == 4); + weights.push_back(*reinterpret_cast(floatWeight)); + if (file.peek() == std::istream::traits_type::eof()) break; + } + std::cout << "Loading weights of " << networkType << " complete" + << std::endl; + std::cout << "Total weights read: " << weights.size() << std::endl; + return weights; +} + +std::string dimsToString(const nvinfer1::Dims d) +{ + std::stringstream s; + assert(d.nbDims >= 1); + for (int i = 0; i < d.nbDims - 1; ++i) + { + s << std::setw(4) << d.d[i] << " x"; + } + s << std::setw(4) << d.d[d.nbDims - 1]; + + return s.str(); +} + +void displayDimType(const nvinfer1::Dims d) +{ + std::cout << "(" << d.nbDims << ") "; + for (int i = 0; i < d.nbDims; ++i) + { + switch (d.type[i]) + { + case nvinfer1::DimensionType::kSPATIAL: std::cout << "kSPATIAL "; break; + case nvinfer1::DimensionType::kCHANNEL: std::cout << "kCHANNEL "; break; + case nvinfer1::DimensionType::kINDEX: std::cout << "kINDEX "; break; + case nvinfer1::DimensionType::kSEQUENCE: std::cout << "kSEQUENCE "; break; + } + } + std::cout << std::endl; +} + +int getNumChannels(nvinfer1::ITensor* t) +{ + nvinfer1::Dims d = t->getDimensions(); + assert(d.nbDims == 3); + + return d.d[0]; +} + +uint64_t get3DTensorVolume(nvinfer1::Dims inputDims) +{ + assert(inputDims.nbDims == 3); + return inputDims.d[0] * inputDims.d[1] * inputDims.d[2]; +} + +void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, + std::string layerOutput, std::string weightPtr) +{ + std::cout << std::setw(6) << std::left << layerIndex << std::setw(24) << std::left << layerName; + std::cout << std::setw(20) << std::left << layerInput << std::setw(20) << std::left + << layerOutput; + std::cout << std::setw(7) << std::left << weightPtr << std::endl; +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/utils.h b/non_square/nvdsinfer_custom_impl_Yolo/utils.h new file mode 100644 index 0000000..2e479f7 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/utils.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + + +#ifndef __UTILS_H__ +#define __UTILS_H__ + +#include +#include +#include +#include +#include + +#include "NvInfer.h" + +#define UNUSED(expr) (void)(expr) +#define DIVUP(n, d) ((n) + (d)-1) / (d) + +std::string trim(std::string s); +float clamp(const float val, const float minVal, const float maxVal); +bool fileExists(const std::string fileName, bool verbose = true); +std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType); +std::string dimsToString(const nvinfer1::Dims d); +void displayDimType(const nvinfer1::Dims d); +int getNumChannels(nvinfer1::ITensor* t); +uint64_t get3DTensorVolume(nvinfer1::Dims inputDims); + +void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, + std::string layerOutput, std::string weightPtr); + +#endif diff --git a/non_square/nvdsinfer_custom_impl_Yolo/yolo.cpp b/non_square/nvdsinfer_custom_impl_Yolo/yolo.cpp new file mode 100644 index 0000000..cc8eb66 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/yolo.cpp @@ -0,0 +1,469 @@ +/* + * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "yolo.h" +#include "yoloPlugins.h" + +void orderParams(std::vector> *maskVector) { + std::vector> maskinput = *maskVector; + std::vector maskPartial; + for (uint i = 0; i < maskinput.size(); i++) { + for (uint j = i + 1; j < maskinput.size(); j++) { + if (maskinput[i][0] <= maskinput[j][0]) { + maskPartial = maskinput[i]; + maskinput[i] = maskinput[j]; + maskinput[j] = maskPartial; + } + } + } + *maskVector = maskinput; +} + +Yolo::Yolo(const NetworkInfo& networkInfo) + : m_NetworkType(networkInfo.networkType), // YOLO type + m_ConfigFilePath(networkInfo.configFilePath), // YOLO cfg + m_WtsFilePath(networkInfo.wtsFilePath), // YOLO weights + m_DeviceType(networkInfo.deviceType), // kDLA, kGPU + m_InputBlobName(networkInfo.inputBlobName), // data + m_InputH(0), + m_InputW(0), + m_InputC(0), + m_InputSize(0) +{} + +Yolo::~Yolo() +{ + destroyNetworkUtils(); +} + +nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder) +{ + assert (builder); + + std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); + std::vector trtWeights; + + nvinfer1::INetworkDefinition *network = builder->createNetwork(); + if (parseModel(*network) != NVDSINFER_SUCCESS) { + network->destroy(); + return nullptr; + } + + // Build the engine + std::cout << "Building the TensorRT Engine" << std::endl; + nvinfer1::ICudaEngine * engine = builder->buildCudaEngine(*network); + if (engine) { + std::cout << "Building complete\n" << std::endl; + } else { + std::cerr << "Building engine failed\n" << std::endl; + } + + // destroy + network->destroy(); + return engine; +} + +NvDsInferStatus Yolo::parseModel(nvinfer1::INetworkDefinition& network) { + destroyNetworkUtils(); + + m_ConfigBlocks = parseConfigFile(m_ConfigFilePath); + parseConfigBlocks(); + orderParams(&m_OutputMasks); + + std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); + // build yolo network + std::cout << "Building YOLO network" << std::endl; + NvDsInferStatus status = buildYoloNetwork(weights, network); + + if (status == NVDSINFER_SUCCESS) { + std::cout << "Building YOLO network complete" << std::endl; + } else { + std::cerr << "Building YOLO network failed" << std::endl; + } + + return status; +} + +NvDsInferStatus Yolo::buildYoloNetwork( + std::vector& weights, nvinfer1::INetworkDefinition& network) { + int weightPtr = 0; + int channels = m_InputC; + + nvinfer1::ITensor* data = + network.addInput(m_InputBlobName.c_str(), nvinfer1::DataType::kFLOAT, + nvinfer1::DimsCHW{static_cast(m_InputC), + static_cast(m_InputH), static_cast(m_InputW)}); + assert(data != nullptr && data->getDimensions().nbDims > 0); + + nvinfer1::ITensor* previous = data; + std::vector tensorOutputs; + uint outputTensorCount = 0; + + // build the network using the network API + for (uint i = 0; i < m_ConfigBlocks.size(); ++i) { + // check if num. of channels is correct + assert(getNumChannels(previous) == channels); + std::string layerIndex = "(" + std::to_string(tensorOutputs.size()) + ")"; + + if (m_ConfigBlocks.at(i).at("type") == "net") { + printLayerInfo("", "layer", " input", " outup", "weightPtr"); + } + + else if (m_ConfigBlocks.at(i).at("type") == "convolutional") { + std::string inputVol = dimsToString(previous->getDimensions()); + nvinfer1::ILayer* out = convolutionalLayer(i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, channels, previous, &network); + previous = out->getOutput(0); + assert(previous != nullptr); + channels = getNumChannels(previous); + std::string outputVol = dimsToString(previous->getDimensions()); + tensorOutputs.push_back(previous); + std::string layerType = "conv_" + m_ConfigBlocks.at(i).at("activation"); + printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr)); + } + + else if (m_ConfigBlocks.at(i).at("type") == "dropout") { + assert(m_ConfigBlocks.at(i).find("probability") != m_ConfigBlocks.at(i).end()); + //float probability = std::stof(m_ConfigBlocks.at(i).at("probability")); + //nvinfer1::ILayer* out = dropoutLayer(probability, previous, &network); + //previous = out->getOutput(0); + //Skip dropout layer + assert(previous != nullptr); + tensorOutputs.push_back(previous); + printLayerInfo(layerIndex, "dropout", " -", " -", " -"); + } + + else if (m_ConfigBlocks.at(i).at("type") == "shortcut") { + assert(m_ConfigBlocks.at(i).find("activation") != m_ConfigBlocks.at(i).end()); + assert(m_ConfigBlocks.at(i).find("from") != m_ConfigBlocks.at(i).end()); + std::string activation = m_ConfigBlocks.at(i).at("activation"); + int from = stoi(m_ConfigBlocks.at(i).at("from")); + if (from > 0) { + from = from - i + 1; + } + assert((i - 2 >= 0) && (i - 2 < tensorOutputs.size())); + assert((i + from - 1 >= 0) && (i + from - 1 < tensorOutputs.size())); + assert(i + from - 1 < i - 2); + std::string inputVol = dimsToString(previous->getDimensions()); + std::string shortcutVol = dimsToString(tensorOutputs[i + from - 1]->getDimensions()); + nvinfer1::ILayer* out = shortcutLayer(i, activation, inputVol, shortcutVol, previous, tensorOutputs[i + from - 1], &network); + previous = out->getOutput(0); + assert(previous != nullptr); + std::string outputVol = dimsToString(previous->getDimensions()); + tensorOutputs.push_back(previous); + std::string layerType = "shortcut_" + m_ConfigBlocks.at(i).at("activation") + ": " + std::to_string(i + from - 1); + printLayerInfo(layerIndex, layerType, " -", outputVol, " -"); + if (inputVol != shortcutVol) { + std::cout << inputVol << " +" << shortcutVol << std::endl; + } + } + + else if (m_ConfigBlocks.at(i).at("type") == "route") { + assert(m_ConfigBlocks.at(i).find("layers") != m_ConfigBlocks.at(i).end()); + nvinfer1::ILayer* out = routeLayer(i, m_ConfigBlocks.at(i), tensorOutputs, &network); + previous = out->getOutput(0); + assert(previous != nullptr); + channels = getNumChannels(previous); + std::string outputVol = dimsToString(previous->getDimensions()); + tensorOutputs.push_back(previous); + printLayerInfo(layerIndex, "route", " -", outputVol, std::to_string(weightPtr)); + } + + else if (m_ConfigBlocks.at(i).at("type") == "upsample") { + std::string inputVol = dimsToString(previous->getDimensions()); + nvinfer1::ILayer* out = upsampleLayer(i - 1, m_ConfigBlocks[i], weights, m_TrtWeights, channels, previous, &network); + previous = out->getOutput(0); + assert(previous != nullptr); + std::string outputVol = dimsToString(previous->getDimensions()); + tensorOutputs.push_back(previous); + printLayerInfo(layerIndex, "upsample", inputVol, outputVol, " -"); + } + + else if (m_ConfigBlocks.at(i).at("type") == "maxpool") { + std::string inputVol = dimsToString(previous->getDimensions()); + nvinfer1::ILayer* out = maxpoolLayer(i, m_ConfigBlocks.at(i), previous, &network); + previous = out->getOutput(0); + assert(previous != nullptr); + std::string outputVol = dimsToString(previous->getDimensions()); + tensorOutputs.push_back(previous); + printLayerInfo(layerIndex, "maxpool", inputVol, outputVol, std::to_string(weightPtr)); + } + + else if (m_ConfigBlocks.at(i).at("type") == "yolo") { + nvinfer1::Dims prevTensorDims = previous->getDimensions(); + assert(prevTensorDims.d[1] == prevTensorDims.d[2]); + TensorInfo& curYoloTensor = m_OutputTensors.at(outputTensorCount); + curYoloTensor.gridSizeY = prevTensorDims.d[1]; + curYoloTensor.gridSizeX = prevTensorDims.d[2]; + curYoloTensor.stride = m_InputH / curYoloTensor.gridSizeY; + m_OutputTensors.at(outputTensorCount).volume = curYoloTensor.gridSizeY + * curYoloTensor.gridSizeX + * (curYoloTensor.numBBoxes * (5 + curYoloTensor.numClasses)); + std::string layerName = "yolo_" + std::to_string(i); + curYoloTensor.blobName = layerName; + int new_coords = 0; + float scale_x_y = 1; + float beta_nms = 0.45; + if (m_ConfigBlocks.at(i).find("new_coords") != m_ConfigBlocks.at(i).end()) { + new_coords = std::stoi(m_ConfigBlocks.at(i).at("new_coords")); + } + if (m_ConfigBlocks.at(i).find("scale_x_y") != m_ConfigBlocks.at(i).end()) { + scale_x_y = std::stof(m_ConfigBlocks.at(i).at("scale_x_y")); + } + if (m_ConfigBlocks.at(i).find("beta_nms") != m_ConfigBlocks.at(i).end()) { + beta_nms = std::stof(m_ConfigBlocks.at(i).at("beta_nms")); + } + nvinfer1::IPluginV2* yoloPlugin + = new YoloLayer(m_OutputTensors.at(outputTensorCount).numBBoxes, + m_OutputTensors.at(outputTensorCount).numClasses, + m_OutputTensors.at(outputTensorCount).gridSizeX, + m_OutputTensors.at(outputTensorCount).gridSizeY, + 1, new_coords, scale_x_y, beta_nms, + curYoloTensor.anchors, + m_OutputMasks); + assert(yoloPlugin != nullptr); + nvinfer1::IPluginV2Layer* yolo = + network.addPluginV2(&previous, 1, *yoloPlugin); + assert(yolo != nullptr); + yolo->setName(layerName.c_str()); + std::string inputVol = dimsToString(previous->getDimensions()); + previous = yolo->getOutput(0); + assert(previous != nullptr); + previous->setName(layerName.c_str()); + std::string outputVol = dimsToString(previous->getDimensions()); + network.markOutput(*previous); + channels = getNumChannels(previous); + tensorOutputs.push_back(yolo->getOutput(0)); + printLayerInfo(layerIndex, "yolo", inputVol, outputVol, std::to_string(weightPtr)); + ++outputTensorCount; + } + + //YOLOv2 support + else if (m_ConfigBlocks.at(i).at("type") == "region") { + nvinfer1::Dims prevTensorDims = previous->getDimensions(); + assert(prevTensorDims.d[1] == prevTensorDims.d[2]); + TensorInfo& curRegionTensor = m_OutputTensors.at(outputTensorCount); + curRegionTensor.gridSizeY = prevTensorDims.d[1]; + curRegionTensor.gridSizeX = prevTensorDims.d[2]; + curRegionTensor.stride = m_InputH / curRegionTensor.gridSizeY; + m_OutputTensors.at(outputTensorCount).volume = curRegionTensor.gridSizeY + * curRegionTensor.gridSizeX + * (curRegionTensor.numBBoxes * (5 + curRegionTensor.numClasses)); + std::string layerName = "region_" + std::to_string(i); + curRegionTensor.blobName = layerName; + std::vector> mask; + nvinfer1::IPluginV2* regionPlugin + = new YoloLayer(curRegionTensor.numBBoxes, + curRegionTensor.numClasses, + curRegionTensor.gridSizeX, + curRegionTensor.gridSizeY, + 0, 0, 1.0, 0, + curRegionTensor.anchors, + mask); + assert(regionPlugin != nullptr); + nvinfer1::IPluginV2Layer* region = + network.addPluginV2(&previous, 1, *regionPlugin); + assert(region != nullptr); + region->setName(layerName.c_str()); + std::string inputVol = dimsToString(previous->getDimensions()); + previous = region->getOutput(0); + assert(previous != nullptr); + previous->setName(layerName.c_str()); + std::string outputVol = dimsToString(previous->getDimensions()); + network.markOutput(*previous); + channels = getNumChannels(previous); + tensorOutputs.push_back(region->getOutput(0)); + printLayerInfo(layerIndex, "region", inputVol, outputVol, std::to_string(weightPtr)); + ++outputTensorCount; + } + else if (m_ConfigBlocks.at(i).at("type") == "reorg") { + std::string inputVol = dimsToString(previous->getDimensions()); + nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2); + assert(reorgPlugin != nullptr); + nvinfer1::IPluginV2Layer* reorg = + network.addPluginV2(&previous, 1, *reorgPlugin); + assert(reorg != nullptr); + std::string layerName = "reorg_" + std::to_string(i); + reorg->setName(layerName.c_str()); + previous = reorg->getOutput(0); + assert(previous != nullptr); + std::string outputVol = dimsToString(previous->getDimensions()); + channels = getNumChannels(previous); + tensorOutputs.push_back(reorg->getOutput(0)); + printLayerInfo(layerIndex, "reorg", inputVol, outputVol, std::to_string(weightPtr)); + } + + else + { + std::cout << "Unsupported layer type --> \"" + << m_ConfigBlocks.at(i).at("type") << "\"" << std::endl; + assert(0); + } + } + + if ((int)weights.size() != weightPtr) + { + std::cout << "Number of unused weights left: " << weights.size() - weightPtr << std::endl; + assert(0); + } + + std::cout << "Output YOLO blob names: " << std::endl; + for (auto& tensor : m_OutputTensors) { + std::cout << tensor.blobName << std::endl; + } + + int nbLayers = network.getNbLayers(); + std::cout << "Total number of YOLO layers: " << nbLayers << std::endl; + + return NVDSINFER_SUCCESS; +} + +std::vector> +Yolo::parseConfigFile (const std::string cfgFilePath) +{ + assert(fileExists(cfgFilePath)); + std::ifstream file(cfgFilePath); + assert(file.good()); + std::string line; + std::vector> blocks; + std::map block; + + while (getline(file, line)) + { + if (line.size() == 0) continue; + if (line.front() == '#') continue; + line = trim(line); + if (line.front() == '[') + { + if (block.size() > 0) + { + blocks.push_back(block); + block.clear(); + } + std::string key = "type"; + std::string value = trim(line.substr(1, line.size() - 2)); + block.insert(std::pair(key, value)); + } + else + { + int cpos = line.find('='); + std::string key = trim(line.substr(0, cpos)); + std::string value = trim(line.substr(cpos + 1)); + block.insert(std::pair(key, value)); + } + } + blocks.push_back(block); + return blocks; +} + +void Yolo::parseConfigBlocks() +{ + for (auto block : m_ConfigBlocks) { + if (block.at("type") == "net") + { + assert((block.find("height") != block.end()) + && "Missing 'height' param in network cfg"); + assert((block.find("width") != block.end()) && "Missing 'width' param in network cfg"); + assert((block.find("channels") != block.end()) + && "Missing 'channels' param in network cfg"); + + m_InputH = std::stoul(block.at("height")); + m_InputW = std::stoul(block.at("width")); + m_InputC = std::stoul(block.at("channels")); + //assert(m_InputW == m_InputH); + m_InputSize = m_InputC * m_InputH * m_InputW; + } + else if ((block.at("type") == "region") || (block.at("type") == "yolo")) + { + assert((block.find("num") != block.end()) + && std::string("Missing 'num' param in " + block.at("type") + " layer").c_str()); + assert((block.find("classes") != block.end()) + && std::string("Missing 'classes' param in " + block.at("type") + " layer") + .c_str()); + assert((block.find("anchors") != block.end()) + && std::string("Missing 'anchors' param in " + block.at("type") + " layer") + .c_str()); + + TensorInfo outputTensor; + std::string anchorString = block.at("anchors"); + while (!anchorString.empty()) + { + int npos = anchorString.find_first_of(','); + if (npos != -1) + { + float anchor = std::stof(trim(anchorString.substr(0, npos))); + outputTensor.anchors.push_back(anchor); + anchorString.erase(0, npos + 1); + } + else + { + float anchor = std::stof(trim(anchorString)); + outputTensor.anchors.push_back(anchor); + break; + } + } + + + if (block.find("mask") != block.end()) { + + std::string maskString = block.at("mask"); + std::vector pMASKS; + while (!maskString.empty()) + { + int npos = maskString.find_first_of(','); + if (npos != -1) + { + int mask = std::stoul(trim(maskString.substr(0, npos))); + pMASKS.push_back(mask); + outputTensor.masks.push_back(mask); + maskString.erase(0, npos + 1); + } + else + { + int mask = std::stoul(trim(maskString)); + pMASKS.push_back(mask); + outputTensor.masks.push_back(mask); + break; + } + } + m_OutputMasks.push_back(pMASKS); + } + + outputTensor.numBBoxes = outputTensor.masks.size() > 0 + ? outputTensor.masks.size() + : std::stoul(trim(block.at("num"))); + outputTensor.numClasses = std::stoul(block.at("classes")); + m_OutputTensors.push_back(outputTensor); + } + } +} + +void Yolo::destroyNetworkUtils() { + // deallocate the weights + for (uint i = 0; i < m_TrtWeights.size(); ++i) { + if (m_TrtWeights[i].count > 0) + free(const_cast(m_TrtWeights[i].values)); + } + m_TrtWeights.clear(); +} \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/yolo.h b/non_square/nvdsinfer_custom_impl_Yolo/yolo.h new file mode 100644 index 0000000..ce1c3be --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/yolo.h @@ -0,0 +1,100 @@ +/* + * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef _YOLO_H_ +#define _YOLO_H_ + +#include "layers/convolutional_layer.h" +#include "layers/dropout_layer.h" +#include "layers/shortcut_layer.h" +#include "layers/route_layer.h" +#include "layers/upsample_layer.h" +#include "layers/maxpool_layer.h" + +#include "nvdsinfer_custom_impl.h" + +struct NetworkInfo +{ + std::string networkType; + std::string configFilePath; + std::string wtsFilePath; + std::string deviceType; + std::string inputBlobName; +}; + +struct TensorInfo +{ + std::string blobName; + uint stride{0}; + uint gridSizeY{0}; + uint gridSizeX{0}; + uint numClasses{0}; + uint numBBoxes{0}; + uint64_t volume{0}; + std::vector masks; + std::vector anchors; + int bindingIndex{-1}; + float* hostBuffer{nullptr}; +}; + +class Yolo : public IModelParser { +public: + Yolo(const NetworkInfo& networkInfo); + ~Yolo() override; + bool hasFullDimsSupported() const override { return false; } + const char* getModelName() const override { + return m_ConfigFilePath.empty() ? m_NetworkType.c_str() + : m_ConfigFilePath.c_str(); + } + NvDsInferStatus parseModel(nvinfer1::INetworkDefinition& network) override; + + nvinfer1::ICudaEngine *createEngine (nvinfer1::IBuilder* builder); + +protected: + const std::string m_NetworkType; + const std::string m_ConfigFilePath; + const std::string m_WtsFilePath; + const std::string m_DeviceType; + const std::string m_InputBlobName; + std::vector m_OutputTensors; + std::vector> m_OutputMasks; + std::vector> m_ConfigBlocks; + uint m_InputH; + uint m_InputW; + uint m_InputC; + uint64_t m_InputSize; + + std::vector m_TrtWeights; + +private: + NvDsInferStatus buildYoloNetwork( + std::vector& weights, nvinfer1::INetworkDefinition& network); + std::vector> parseConfigFile( + const std::string cfgFilePath); + void parseConfigBlocks(); + void destroyNetworkUtils(); +}; + +#endif // _YOLO_H_ \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/yoloForward.cu b/non_square/nvdsinfer_custom_impl_Yolo/yoloForward.cu new file mode 100644 index 0000000..dcc4b95 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/yoloForward.cu @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2018-2019 NVIDIA Corporation. All rights reserved. + * + * NVIDIA Corporation and its licensors retain all intellectual property + * and proprietary rights in and to this software, related documentation + * and any modifications thereto. Any use, reproduction, disclosure or + * distribution of this software and related documentation without an express + * license agreement from NVIDIA Corporation is strictly prohibited. + * + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + * + */ + +#include +#include +#include +#include +#include + +inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } + +__global__ void gpuYoloLayer(const float* input, float* output, const uint gridSizeX, const uint gridSizeY, const uint numOutputClasses, + const uint numBBoxes, const uint new_coords, const float scale_x_y) +{ + uint x_id = blockIdx.x * blockDim.x + threadIdx.x; + uint y_id = blockIdx.y * blockDim.y + threadIdx.y; + uint z_id = blockIdx.z * blockDim.z + threadIdx.z; + + if ((x_id >= gridSizeX) || (y_id >= gridSizeY) || (z_id >= numBBoxes)) + { + return; + } + + const int numGridCells = gridSizeX * gridSizeY; + const int bbindex = y_id * gridSizeX + x_id; + + float alpha = scale_x_y; + float beta = -0.5 * (scale_x_y - 1); + + if (new_coords == 1) { + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] + = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] * alpha + beta; + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] + = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] * alpha + beta; + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] + = pow(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] * 2, 2); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] + = pow(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] * 2, 2); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] + = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]; + + for (uint i = 0; i < numOutputClasses; ++i) + { + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] + = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]; + } + } + else { + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]) * alpha + beta; + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]) * alpha + beta; + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] + = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] + = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); + + for (uint i = 0; i < numOutputClasses; ++i) + { + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]); + } + } +} + +__global__ void gpuRegionLayer(const float* input, float* output, const uint gridSizeX, const uint gridSizeY, const uint numOutputClasses, + const uint numBBoxes) +{ + uint x_id = blockIdx.x * blockDim.x + threadIdx.x; + uint y_id = blockIdx.y * blockDim.y + threadIdx.y; + uint z_id = blockIdx.z * blockDim.z + threadIdx.z; + + if ((x_id >= gridSizeX) || (y_id >= gridSizeY) || (z_id >= numBBoxes)) + { + return; + } + + const int numGridCells = gridSizeX * gridSizeY; + const int bbindex = y_id * gridSizeX + x_id; + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] + = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] + = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]); + + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] + = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); + + float temp = 1.0; + int i; + float sum = 0; + float largest = -INFINITY; + for(i = 0; i < numOutputClasses; ++i){ + int val = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]; + largest = (val>largest) ? val : largest; + } + for(i = 0; i < numOutputClasses; ++i){ + float e = exp(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] / temp - largest / temp); + sum += e; + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] = e; + } + for(i = 0; i < numOutputClasses; ++i){ + output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] /= sum; + } +} + +cudaError_t cudaYoloLayer(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, + const uint& numOutputClasses, const uint& numBBoxes, + uint64_t outputSize, cudaStream_t stream, const uint modelCoords, const float modelScale, const uint modelType); + +cudaError_t cudaYoloLayer(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, + const uint& numOutputClasses, const uint& numBBoxes, + uint64_t outputSize, cudaStream_t stream, const uint modelCoords, const float modelScale, const uint modelType) +{ + dim3 threads_per_block(16, 16, 4); + dim3 number_of_blocks((gridSizeX / threads_per_block.x) + 1, + (gridSizeY / threads_per_block.y) + 1, + (numBBoxes / threads_per_block.z) + 1); + if (modelType == 1) { + for (unsigned int batch = 0; batch < batchSize; ++batch) + { + gpuYoloLayer<<>>( + reinterpret_cast(input) + (batch * outputSize), + reinterpret_cast(output) + (batch * outputSize), gridSizeX, gridSizeY, numOutputClasses, + numBBoxes, modelCoords, modelScale); + } + } + else if (modelType == 0) { + for (unsigned int batch = 0; batch < batchSize; ++batch) + { + gpuRegionLayer<<>>( + reinterpret_cast(input) + (batch * outputSize), + reinterpret_cast(output) + (batch * outputSize), gridSizeX, gridSizeY, numOutputClasses, + numBBoxes); + } + } + return cudaGetLastError(); +} diff --git a/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp b/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp new file mode 100644 index 0000000..0ae7cbb --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp @@ -0,0 +1,209 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include "yoloPlugins.h" +#include "NvInferPlugin.h" +#include +#include +#include + +int kNUM_CLASSES; +float kBETA_NMS; +std::vector kANCHORS; +std::vector> kMASK; + +namespace { +template +void write(char*& buffer, const T& val) +{ + *reinterpret_cast(buffer) = val; + buffer += sizeof(T); +} + +template +void read(const char*& buffer, T& val) +{ + val = *reinterpret_cast(buffer); + buffer += sizeof(T); +} +} + +cudaError_t cudaYoloLayer ( + const void* input, void* output, const uint& batchSize, + const uint& gridSizeX, const uint& gridSizeY, const uint& numOutputClasses, + const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, const uint modelCoords, const float modelScale, const uint modelType); + +YoloLayer::YoloLayer (const void* data, size_t length) +{ + const char *d = static_cast(data); + read(d, m_NumBoxes); + read(d, m_NumClasses); + read(d, m_GridSizeX); + read(d, m_GridSizeY); + read(d, m_OutputSize); + + read(d, m_type); + read(d, m_new_coords); + read(d, m_scale_x_y); + read(d, m_beta_nms); + uint anchorsSize; + read(d, anchorsSize); + for (uint i = 0; i < anchorsSize; i++) { + float result; + read(d, result); + m_Anchors.push_back(result); + } + uint maskSize; + read(d, maskSize); + for (uint i = 0; i < maskSize; i++) { + uint nMask; + read(d, nMask); + std::vector pMask; + for (uint f = 0; f < nMask; f++) { + int result; + read(d, result); + pMask.push_back(result); + } + m_Mask.push_back(pMask); + } + kNUM_CLASSES = m_NumClasses; + kBETA_NMS = m_beta_nms; + kANCHORS = m_Anchors; + kMASK = m_Mask; +}; + +YoloLayer::YoloLayer ( + const uint& numBoxes, const uint& numClasses, const uint& gridSizeX, const uint& gridSizeY, const uint model_type, const uint new_coords, const float scale_x_y, const float beta_nms, const std::vector anchors, std::vector> mask) : + m_NumBoxes(numBoxes), + m_NumClasses(numClasses), + m_GridSizeX(gridSizeX), + m_GridSizeY(gridSizeY), + m_type(model_type), + m_new_coords(new_coords), + m_scale_x_y(scale_x_y), + m_beta_nms(beta_nms), + m_Anchors(anchors), + m_Mask(mask) +{ + assert(m_NumBoxes > 0); + assert(m_NumClasses > 0); + assert(m_GridSizeX > 0); + assert(m_GridSizeY > 0); + m_OutputSize = m_GridSizeX * m_GridSizeY * (m_NumBoxes * (4 + 1 + m_NumClasses)); +}; + +nvinfer1::Dims +YoloLayer::getOutputDimensions( + int index, const nvinfer1::Dims* inputs, int nbInputDims) +{ + assert(index == 0); + assert(nbInputDims == 1); + return inputs[0]; +} + +bool YoloLayer::supportsFormat ( + nvinfer1::DataType type, nvinfer1::PluginFormat format) const { + return (type == nvinfer1::DataType::kFLOAT && + format == nvinfer1::PluginFormat::kNCHW); +} + +void +YoloLayer::configureWithFormat ( + const nvinfer1::Dims* inputDims, int nbInputs, + const nvinfer1::Dims* outputDims, int nbOutputs, + nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) +{ + assert(nbInputs == 1); + assert (format == nvinfer1::PluginFormat::kNCHW); + assert(inputDims != nullptr); +} + +int YoloLayer::enqueue( + int batchSize, const void* const* inputs, void** outputs, void* workspace, + cudaStream_t stream) +{ + CHECK(cudaYoloLayer( + inputs[0], outputs[0], batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, + m_OutputSize, stream, m_new_coords, m_scale_x_y, m_type)); + return 0; +} + +size_t YoloLayer::getSerializationSize() const +{ + int anchorsSum = 1; + for (uint i = 0; i < m_Anchors.size(); i++) { + anchorsSum += 1; + } + int maskSum = 1; + for (uint i = 0; i < m_Mask.size(); i++) { + maskSum += 1; + for (uint f = 0; f < m_Mask[i].size(); f++) { + maskSum += 1; + } + } + + return sizeof(m_NumBoxes) + sizeof(m_NumClasses) + sizeof(m_GridSizeX) + sizeof(m_GridSizeY) + sizeof(m_OutputSize) + sizeof(m_type) + + sizeof(m_new_coords) + sizeof(m_scale_x_y) + sizeof(m_beta_nms) + anchorsSum * sizeof(float) + maskSum * sizeof(int); +} + +void YoloLayer::serialize(void* buffer) const +{ + char *d = static_cast(buffer); + write(d, m_NumBoxes); + write(d, m_NumClasses); + write(d, m_GridSizeX); + write(d, m_GridSizeY); + write(d, m_OutputSize); + + write(d, m_type); + write(d, m_new_coords); + write(d, m_scale_x_y); + write(d, m_beta_nms); + uint anchorsSize = m_Anchors.size(); + write(d, anchorsSize); + for (uint i = 0; i < anchorsSize; i++) { + write(d, m_Anchors[i]); + } + uint maskSize = m_Mask.size(); + write(d, maskSize); + for (uint i = 0; i < maskSize; i++) { + uint pMaskSize = m_Mask[i].size(); + write(d, pMaskSize); + for (uint f = 0; f < pMaskSize; f++) { + write(d, m_Mask[i][f]); + } + } + kNUM_CLASSES = m_NumClasses; + kBETA_NMS = m_beta_nms; + kANCHORS = m_Anchors; + kMASK = m_Mask; +} + +nvinfer1::IPluginV2* YoloLayer::clone() const +{ + return new YoloLayer (m_NumBoxes, m_NumClasses, m_GridSizeX, m_GridSizeY, m_type, m_new_coords, m_scale_x_y, m_beta_nms, m_Anchors, m_Mask); +} + +REGISTER_TENSORRT_PLUGIN(YoloLayerPluginCreator); \ No newline at end of file diff --git a/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h b/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h new file mode 100644 index 0000000..2981af7 --- /dev/null +++ b/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h @@ -0,0 +1,156 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + + * Edited by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#ifndef __YOLO_PLUGINS__ +#define __YOLO_PLUGINS__ + +#include +#include +#include +#include +#include + +#include + +#include "NvInferPlugin.h" + +#define CHECK(status) \ + { \ + if (status != 0) \ + { \ + std::cout << "CUDA failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \ + << " at line " << __LINE__ << std::endl; \ + abort(); \ + } \ + } + +namespace +{ +const char* YOLOLAYER_PLUGIN_VERSION {"1"}; +const char* YOLOLAYER_PLUGIN_NAME {"YoloLayer_TRT"}; +} // namespace + +class YoloLayer : public nvinfer1::IPluginV2 +{ +public: + YoloLayer (const void* data, size_t length); + YoloLayer (const uint& numBoxes, const uint& numClasses, const uint& gridSize, + const uint model_type, const uint new_coords, const float scale_x_y, const float beta_nms, + const std::vector anchors, const std::vector> mask); + const char* getPluginType () const override { return YOLOLAYER_PLUGIN_NAME; } + const char* getPluginVersion () const override { return YOLOLAYER_PLUGIN_VERSION; } + int getNbOutputs () const override { return 1; } + + nvinfer1::Dims getOutputDimensions ( + int index, const nvinfer1::Dims* inputs, + int nbInputDims) override; + + bool supportsFormat ( + nvinfer1::DataType type, nvinfer1::PluginFormat format) const override; + + void configureWithFormat ( + const nvinfer1::Dims* inputDims, int nbInputs, + const nvinfer1::Dims* outputDims, int nbOutputs, + nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) override; + + int initialize () override { return 0; } + void terminate () override {} + size_t getWorkspaceSize (int maxBatchSize) const override { return 0; } + int enqueue ( + int batchSize, const void* const* inputs, void** outputs, + void* workspace, cudaStream_t stream) override; + size_t getSerializationSize() const override; + void serialize (void* buffer) const override; + void destroy () override { delete this; } + nvinfer1::IPluginV2* clone() const override; + + void setPluginNamespace (const char* pluginNamespace)override { + m_Namespace = pluginNamespace; + } + virtual const char* getPluginNamespace () const override { + return m_Namespace.c_str(); + } + +private: + uint m_NumBoxes {0}; + uint m_NumClasses {0}; + uint m_GridSizeX {0}; + uint m_GridSizeY {0}; + uint64_t m_OutputSize {0}; + std::string m_Namespace {""}; + + uint m_type {0}; + uint m_new_coords {0}; + float m_scale_x_y {0}; + float m_beta_nms {0}; + std::vector m_Anchors; + std::vector> m_Mask; +}; + +class YoloLayerPluginCreator : public nvinfer1::IPluginCreator +{ +public: + YoloLayerPluginCreator () {} + ~YoloLayerPluginCreator () {} + + const char* getPluginName () const override { return YOLOLAYER_PLUGIN_NAME; } + const char* getPluginVersion () const override { return YOLOLAYER_PLUGIN_VERSION; } + + const nvinfer1::PluginFieldCollection* getFieldNames() override { + std::cerr<< "YoloLayerPluginCreator::getFieldNames is not implemented" << std::endl; + return nullptr; + } + + nvinfer1::IPluginV2* createPlugin ( + const char* name, const nvinfer1::PluginFieldCollection* fc) override + { + std::cerr<< "YoloLayerPluginCreator::getFieldNames is not implemented"; + return nullptr; + } + + nvinfer1::IPluginV2* deserializePlugin ( + const char* name, const void* serialData, size_t serialLength) override + { + std::cout << "Deserialize yoloLayer plugin: " << name << std::endl; + return new YoloLayer(serialData, serialLength); + } + + void setPluginNamespace(const char* libNamespace) override { + m_Namespace = libNamespace; + } + const char* getPluginNamespace() const override { + return m_Namespace.c_str(); + } + +private: + std::string m_Namespace {""}; +}; + +extern int kNUM_CLASSES; +extern float kBETA_NMS; +extern std::vector kANCHORS; +extern std::vector> kMASK; + +#endif // __YOLO_PLUGINS__ \ No newline at end of file