Add YOLOv8 support

2023-01-27 15:56:00 -03:00
parent f1cd701247
commit f9c7a4dfca
59 changed files with 3260 additions and 2763 deletions
--- a/nvdsinfer_custom_impl_Yolo/layers/c2f_layer.cpp
+++ b/nvdsinfer_custom_impl_Yolo/layers/c2f_layer.cpp
@@ -0,0 +1,82 @@
+/*
+ * Created by Marcos Luciano
+ * https://www.github.com/marcoslucianops
+ */
+
+#include "c2f_layer.h"
+
+#include <cassert>
+
+#include "convolutional_layer.h"
+
+nvinfer1::ITensor*
+c2fLayer(int layerIdx, std::map<std::string, std::string>& block, std::vector<float>& weights,
+    std::vector<nvinfer1::Weights>& trtWeights, int& weightPtr, std::string weightsType, float eps, nvinfer1::ITensor* input,
+    nvinfer1::INetworkDefinition* network)
+{
+  nvinfer1::ITensor* output;
+
+  assert(block.at("type") == "c2f");
+  assert(block.find("n") != block.end());
+  assert(block.find("shortcut") != block.end());
+  assert(block.find("filters") != block.end());
+
+  int n = std::stoi(block.at("n"));
+  bool shortcut = (block.at("shortcut") == "1");
+  int filters = std::stoi(block.at("filters"));
+
+  nvinfer1::Dims inputDims = input->getDimensions();
+
+  nvinfer1::ISliceLayer* sliceLt = network->addSlice(*input,nvinfer1::Dims{3, {0, 0, 0}},
+      nvinfer1::Dims{3, {inputDims.d[0] / 2, inputDims.d[1], inputDims.d[2]}}, nvinfer1::Dims{3, {1, 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(*input,nvinfer1::Dims{3, {inputDims.d[0] / 2, 0, 0}},
+      nvinfer1::Dims{3, {inputDims.d[0] / 2, inputDims.d[1], inputDims.d[2]}}, nvinfer1::Dims{3, {1, 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);
+
+  std::vector<nvinfer1::ITensor*> concatInputs;
+  concatInputs.push_back(lt);
+  concatInputs.push_back(rb);
+  output = rb;
+
+  for (int i = 0; i < n; ++i) {
+    std::string cv1MlayerName = "c2f_1_" + std::to_string(i + 1) + "_";
+    nvinfer1::ITensor* cv1M = convolutionalLayer(layerIdx, block, weights, trtWeights, weightPtr, weightsType, filters, eps,
+        output, network, cv1MlayerName);
+    assert(cv1M != nullptr);
+
+    std::string cv2MlayerName = "c2f_2_" + std::to_string(i + 1) + "_";
+    nvinfer1::ITensor* cv2M = convolutionalLayer(layerIdx, block, weights, trtWeights, weightPtr, weightsType, filters, eps,
+        cv1M, network, cv2MlayerName);
+    assert(cv2M != nullptr);
+
+    if (shortcut) {
+      nvinfer1::IElementWiseLayer* ew = network->addElementWise(*rb, *cv2M, nvinfer1::ElementWiseOperation::kSUM);
+      assert(ew != nullptr);
+      std::string ewLayerName = "shortcut_c2f_" + std::to_string(i + 1) + "_" + std::to_string(layerIdx);
+      ew->setName(ewLayerName.c_str());
+      output = ew->getOutput(0);
+      concatInputs.push_back(output);
+    }
+    else {
+      output = cv2M;
+      concatInputs.push_back(output);
+    }
+  }
+
+  nvinfer1::IConcatenationLayer* concat = network->addConcatenation(concatInputs.data(), concatInputs.size());
+  assert(concat != nullptr);
+  std::string concatLayerName = "route_" + std::to_string(layerIdx);
+  concat->setName(concatLayerName.c_str());
+  concat->setAxis(0);
+  output = concat->getOutput(0);
+
+  return output;
+}