Add PP-YOLOE support

nvdsinfer_custom_impl_Yolo/layers/reg_layer.cpp

@@ -0,0 +1,113 @@
+/*
+ * Created by Marcos Luciano
+ * https://www.github.com/marcoslucianops
+ */
+
+#include "reg_layer.h"
+
+nvinfer1::ITensor* regLayer(
+    int layerIdx,
+    std::map<std::string, std::string>& block,
+    std::vector<float>& weights,
+    std::vector<nvinfer1::Weights>& trtWeights,
+    int& weightPtr,
+    nvinfer1::ITensor* input,
+    nvinfer1::INetworkDefinition* network)
+{
+    nvinfer1::ITensor* output;
+
+    assert(block.at("type") == "reg");
+
+    nvinfer1::IShuffleLayer* shuffle = network->addShuffle(*input);
+    assert(shuffle != nullptr);
+    std::string shuffleLayerName = "shuffle_" + std::to_string(layerIdx);
+    shuffle->setName(shuffleLayerName.c_str());
+    nvinfer1::Permutation permutation;
+    permutation.order[0] = 1;
+    permutation.order[1] = 0;
+    shuffle->setFirstTranspose(permutation);
+    output = shuffle->getOutput(0);
+    nvinfer1::Dims shuffleDims = output->getDimensions();
+
+    nvinfer1::ISliceLayer* sliceLt = network->addSlice(
+        *output, nvinfer1::Dims{2, {0, 0}}, nvinfer1::Dims{2, {shuffleDims.d[0], 2}}, nvinfer1::Dims{2, {1, 1}});
+    assert(sliceLt != nullptr);
+    std::string sliceLtLayerName = "slice_lt_" + std::to_string(layerIdx);
+    sliceLt->setName(sliceLtLayerName.c_str());
+    nvinfer1::ITensor* lt = sliceLt->getOutput(0);
+
+    nvinfer1::ISliceLayer* sliceRb = network->addSlice(
+        *output, nvinfer1::Dims{2, {0, 2}}, nvinfer1::Dims{2, {shuffleDims.d[0], 2}}, nvinfer1::Dims{2, {1, 1}});
+    assert(sliceRb != nullptr);
+    std::string sliceRbLayerName = "slice_rb_" + std::to_string(layerIdx);
+    sliceRb->setName(sliceRbLayerName.c_str());
+    nvinfer1::ITensor* rb = sliceRb->getOutput(0);
+
+    int channels = shuffleDims.d[0] * 2;
+    nvinfer1::Weights anchorPointsWt{nvinfer1::DataType::kFLOAT, nullptr, channels};
+    float* val = new float[channels];
+    for (int i = 0; i < channels; ++i)
+    {
+        val[i] = weights[weightPtr];
+        weightPtr++;
+    }
+    anchorPointsWt.values = val;
+    trtWeights.push_back(anchorPointsWt);
+
+    nvinfer1::IConstantLayer* anchorPoints = network->addConstant(nvinfer1::Dims{2, {shuffleDims.d[0], 2}}, anchorPointsWt);
+    assert(anchorPoints != nullptr);
+    std::string anchorPointsLayerName = "anchor_points_" + std::to_string(layerIdx);
+    anchorPoints->setName(anchorPointsLayerName.c_str());
+    nvinfer1::ITensor* anchorPointsTensor = anchorPoints->getOutput(0);
+
+    nvinfer1::IElementWiseLayer* x1y1
+        = network->addElementWise(*anchorPointsTensor, *lt, nvinfer1::ElementWiseOperation::kSUB);
+    assert(x1y1 != nullptr);
+    std::string x1y1LayerName = "x1y1_" + std::to_string(layerIdx);
+    x1y1->setName(x1y1LayerName.c_str());
+    nvinfer1::ITensor* x1y1Tensor = x1y1->getOutput(0);
+
+    nvinfer1::IElementWiseLayer* x2y2
+        = network->addElementWise(*rb, *anchorPointsTensor, nvinfer1::ElementWiseOperation::kSUM);
+    assert(x2y2 != nullptr);
+    std::string x2y2LayerName = "x2y2_" + std::to_string(layerIdx);
+    x2y2->setName(x2y2LayerName.c_str());
+    nvinfer1::ITensor* x2y2Tensor = x2y2->getOutput(0);
+
+    std::vector<nvinfer1::ITensor*> concatInputs;
+    concatInputs.push_back(x1y1Tensor);
+    concatInputs.push_back(x2y2Tensor);
+
+    nvinfer1::IConcatenationLayer* concat = network->addConcatenation(concatInputs.data(), concatInputs.size());
+    assert(concat != nullptr);
+    std::string concatLayerName = "concat_" + std::to_string(layerIdx);
+    concat->setName(concatLayerName.c_str());
+    concat->setAxis(1);
+    output = concat->getOutput(0);
+
+    channels = shuffleDims.d[0];
+    nvinfer1::Weights stridePointsWt{nvinfer1::DataType::kFLOAT, nullptr, channels};
+    val = new float[channels];
+    for (int i = 0; i < channels; ++i)
+    {
+        val[i] = weights[weightPtr];
+        weightPtr++;
+    }
+    stridePointsWt.values = val;
+    trtWeights.push_back(stridePointsWt);
+
+    nvinfer1::IConstantLayer* stridePoints = network->addConstant(nvinfer1::Dims{2, {shuffleDims.d[0], 1}}, stridePointsWt);
+    assert(stridePoints != nullptr);
+    std::string stridePointsLayerName = "stride_points_" + std::to_string(layerIdx);
+    stridePoints->setName(stridePointsLayerName.c_str());
+    nvinfer1::ITensor* stridePointsTensor = stridePoints->getOutput(0);
+
+    nvinfer1::IElementWiseLayer* pred
+        = network->addElementWise(*output, *stridePointsTensor, nvinfer1::ElementWiseOperation::kPROD);
+    assert(pred != nullptr);
+    std::string predLayerName = "pred_" + std::to_string(layerIdx);
+    pred->setName(predLayerName.c_str());
+    output = pred->getOutput(0);
+
+    return output;
+}