separate model and block (wang-xinyu#1154)

Author: wang-xinyu

yolov7/include/block.h

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "NvInfer.h"
+#include <string>
+#include <vector>
+#include <map>
+
+std::map<std::string, nvinfer1::Weights> loadWeights(const std::string file);
+
+nvinfer1::IElementWiseLayer* convBnSilu(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c2, int k, int s, int p, std::string lname);
+
+nvinfer1::ILayer* ReOrg(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int inch);
+
+nvinfer1::ILayer* DownC(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c1, int c2, const std::string& lname);
+
+nvinfer1::IElementWiseLayer* SPPCSPC(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c2, const std::string& lname);
+
+nvinfer1::IElementWiseLayer* RepConv(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c2, int k, int s, const std::string& lname);
+
+nvinfer1::IActivationLayer* convBlockLeakRelu(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int outch, int ksize, int s, int p, std::string lname);
+
+nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::map<std::string, nvinfer1::Weights>& weightMap, std::string lname, std::vector<nvinfer1::IConvolutionLayer*> dets);
+

yolov7/include/common.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
+
+// These are used to define input/output tensor names,
+// you can set them to whatever you want.
+const static char* kInputTensorName = "data";
+const static char* kOutputTensorName = "prob";
+
+const static int kNumClass = 80;
+
+// Yolo's input width and height must by divisible by 32
+const static int kInputH = 640;
+const static int kInputW = 640;
+
+// Maximum number of output bounding boxes from yololayer plugin.
+// That is maximum number of output bounding boxes before NMS.
+const static int kMaxNumOutputBbox = 1000;
+
+const static int kNumAnchor = 3;
+
+// The bboxes whose conf lower kIgnoreThresh will be ignored in yololayer plugin.
+const static float kIgnoreThresh = 0.1f;
+

yolov7/include/config.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include "NvInfer.h"
+#include <string>
+
+nvinfer1::ICudaEngine* build_engine_yolov7e6e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::ICudaEngine* build_engine_yolov7d6(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::ICudaEngine* build_engine_yolov7e6(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::ICudaEngine* build_engine_yolov7w6(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::ICudaEngine* build_engine_yolov7x(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::ICudaEngine* build_engine_yolov7(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::ICudaEngine* build_engine_yolov7_tiny(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name);

yolov7/include/model.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "types.h"
+#include <opencv2/opencv.hpp>
+
+cv::Rect get_rect(cv::Mat& img, float bbox[4]);
+
+void nms(std::vector<Detection>& res, float *output, float conf_thresh, float nms_thresh = 0.5);
+

yolov7/include/postprocess.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "config.h"
+
+struct YoloKernel {
+  int width;
+  int height;
+  float anchors[kNumAnchor * 2];
+};
+
+struct alignas(float) Detection {
+  //center_x center_y w h
+  float bbox[4];
+  float conf;  // bbox_conf * cls_conf
+  float class_id;
+};
+

yolov7/include/types.h

@@ -19,7 +19,7 @@ void read(const char*& buffer, T& val) {
 }  // namespace Tn
 
 namespace nvinfer1 {
-YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<Yolo::YoloKernel>& vYoloKernel) {
+YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<YoloKernel>& vYoloKernel) {
   mClassCount = classCount;
   mYoloV7NetWidth = netWidth;
   mYoloV7NetHeight = netHeight;
@@ -28,7 +28,7 @@ YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, in
   mKernelCount = vYoloKernel.size();
 
   CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
-  size_t AnchorLen = sizeof(float) * Yolo::CHECK_COUNT * 2;
+  size_t AnchorLen = sizeof(float) * kNumAnchor * 2;
   for (int ii = 0; ii < mKernelCount; ii++) {
     CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen));
     const auto& yolo = mYoloKernel[ii];
@@ -53,11 +53,11 @@ YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) {
   read(d, mYoloV7NetHeight);
   read(d, mMaxOutObject);
   mYoloKernel.resize(mKernelCount);
-  auto kernelSize = mKernelCount * sizeof(Yolo::YoloKernel);
+  auto kernelSize = mKernelCount * sizeof(YoloKernel);
   memcpy(mYoloKernel.data(), d, kernelSize);
   d += kernelSize;
   CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*)));
-  size_t AnchorLen = sizeof(float) * Yolo::CHECK_COUNT * 2;
+  size_t AnchorLen = sizeof(float) * kNumAnchor * 2;
   for (int ii = 0; ii < mKernelCount; ii++) {
     CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen));
     const auto& yolo = mYoloKernel[ii];
@@ -75,15 +75,15 @@ void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT {
   write(d, mYoloV7NetWidth);
   write(d, mYoloV7NetHeight);
   write(d, mMaxOutObject);
-  auto kernelSize = mKernelCount * sizeof(Yolo::YoloKernel);
+  auto kernelSize = mKernelCount * sizeof(YoloKernel);
   memcpy(d, mYoloKernel.data(), kernelSize);
   d += kernelSize;
 
   assert(d == a + getSerializationSize());
 }
 
 size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT {
-  return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV7NetWidth) + sizeof(mYoloV7NetHeight) + sizeof(mMaxOutObject);
+  return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(YoloKernel) * mYoloKernel.size() + sizeof(mYoloV7NetWidth) + sizeof(mYoloV7NetHeight) + sizeof(mMaxOutObject);
 }
 
 int YoloLayerPlugin::initialize() TRT_NOEXCEPT {
@@ -92,7 +92,7 @@ int YoloLayerPlugin::initialize() TRT_NOEXCEPT {
 
 Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT {
   //output the result to channel
-  int totalsize = mMaxOutObject * sizeof(Yolo::Detection) / sizeof(float);
+  int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float);
   return Dims3(totalsize + 1, 1, 1);
 }
 
@@ -150,19 +150,19 @@ IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT {
 __device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); };
 
 __global__ void CalDetection(const float *input, float *output, int noElements,
-    const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[Yolo::CHECK_COUNT * 2], int classes, int outputElem) {
+    const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[kNumAnchor * 2], int classes, int outputElem) {
   int idx = threadIdx.x + blockDim.x * blockIdx.x;
   if (idx >= noElements) return;
 
   int total_grid = yoloWidth * yoloHeight;  // 80*80 40*40 20*20
   int bnIdx = idx / total_grid;
   idx = idx - total_grid * bnIdx;
   int info_len_i = 5 + classes;
-  const float* curInput = input + bnIdx * (info_len_i * total_grid * Yolo::CHECK_COUNT);
+  const float* curInput = input + bnIdx * (info_len_i * total_grid * kNumAnchor);
 
   for (int k = 0; k < 3; k++) {
     float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]);
-    if (box_prob < Yolo::IGNORE_THRESH) continue;
+    if (box_prob < kIgnoreThresh) continue;
     int class_id = 0;
     float max_cls_prob = 0.0;
     for (int i = 5; i < info_len_i; ++i) {
@@ -175,8 +175,8 @@ __global__ void CalDetection(const float *input, float *output, int noElements,
     float *res_count = output + bnIdx * outputElem;
     int count = (int)atomicAdd(res_count, 1);
     if (count >= maxoutobject) return;
-    char *data = (char*)res_count + sizeof(float) + count * sizeof(Yolo::Detection);
-    Yolo::Detection *det = (Yolo::Detection*)(data);
+    char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection);
+    Detection *det = (Detection*)(data);
 
     int row = idx / yoloWidth;
     int col = idx % yoloWidth;
@@ -204,7 +204,7 @@ __global__ void CalDetection(const float *input, float *output, int noElements,
 }
 
 void YoloLayerPlugin::forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize) {
-  int outputElem = 1 + mMaxOutObject * sizeof(Yolo::Detection) / sizeof(float);
+  int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);
   for (int idx = 0; idx < batchSize; ++idx) {
     CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream));
   }
@@ -255,8 +255,8 @@ IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFi
   int input_w = p_netinfo[1];
   int input_h = p_netinfo[2];
   int max_output_object_count = p_netinfo[3];
-  std::vector<Yolo::YoloKernel> kernels(fc->fields[1].length);
-  memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(Yolo::YoloKernel));
+  std::vector<YoloKernel> kernels(fc->fields[1].length);
+  memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(YoloKernel));
   YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, kernels);
   obj->setPluginNamespace(mNamespace.c_str());
   return obj;

yolov7/main.cpp

@@ -1,36 +1,14 @@
-#ifndef _YOLO_LAYER_H
-#define _YOLO_LAYER_H
+#pragma once
 
 #include "macros.h"
+#include "types.h"
 #include <vector>
 #include <string>
 
-namespace Yolo {
-static constexpr int CHECK_COUNT = 3;
-static constexpr float IGNORE_THRESH = 0.5f;
-struct YoloKernel {
-  int width;
-  int height;
-  float anchors[CHECK_COUNT * 2];
-};
-static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000;
-static constexpr int CLASS_NUM = 80;
-static constexpr int INPUT_H = 640;  // yolov7's input height and width must be divisible by 32.
-static constexpr int INPUT_W = 640;
-
-static constexpr int LOCATIONS = 4;
-struct alignas(float) Detection {
-  //center_x center_y w h
-  float bbox[LOCATIONS];
-  float conf;  // bbox_conf * cls_conf
-  float class_id;
-};
-}  // namespace yolo
-
 namespace nvinfer1 {
 class API YoloLayerPlugin : public IPluginV2IOExt {
  public:
-  YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<Yolo::YoloKernel>& vYoloKernel);
+  YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<YoloKernel>& vYoloKernel);
   YoloLayerPlugin(const void* data, size_t length);
   ~YoloLayerPlugin();
 
@@ -90,7 +68,7 @@ class API YoloLayerPlugin : public IPluginV2IOExt {
   int mYoloV7NetWidth;
   int mYoloV7NetHeight;
   int mMaxOutObject;
-  std::vector<Yolo::YoloKernel> mYoloKernel;
+  std::vector<YoloKernel> mYoloKernel;
   void** mAnchor;
 };
 
@@ -126,5 +104,3 @@ class API YoloPluginCreator : public IPluginCreator {
 REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
 }  // namespace nvinfer1
 
-#endif  // _YOLO_LAYER_H
-