yolov8 detection p6 (#1479)

* yolov8 detection p6 * Update README.md * update README.md * Update README.md * Update README.md * Update README.md * Update README.md * update README.md
wang-xinyu · Apr 1, 2024 · 62050de · 62050de
1 parent c86d808
commit 62050de
Show file tree

Hide file tree

Showing 6 changed files with 818 additions and 290 deletions.
diff --git a/yolov8/README.md b/yolov8/README.md
@@ -19,13 +19,13 @@ The tensorrt code is derived from [xiaocao-tian/yolov8_tensorrt](https://github.
 
 ## Different versions of yolov8
 
-Currently, we support yolov8 
+Currently, we support yolov8
 
 - For yolov8 , download .pt from [https://github.com/ultralytics/assets/releases](https://github.com/ultralytics/assets/releases), then follow how-to-run in current page.
 
 ## Config
 
-- Choose the model n/s/m/l/x from command line arguments.
+- Choose the model n/s/m/l/x/n6/s6/m6/l6/x6 from command line arguments.
 - Check more configs in [include/config.h](./include/config.h)
 
 ## How to Run, yolov8n as example
@@ -51,7 +51,7 @@ cd build
 cp {ultralytics}/ultralytics/yolov8.wts {tensorrtx}/yolov8/build
 cmake ..
 make
-sudo ./yolov8_det -s [.wts] [.engine] [n/s/m/l/x]  // serialize model to plan file
+sudo ./yolov8_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6]  // serialize model to plan file
 sudo ./yolov8_det -d [.engine] [image folder]  [c/g] // deserialize and run inference, the images in [image folder] will be processed.
 // For example yolov8
 sudo ./yolov8_det -s yolov8n.wts yolov8.engine n
@@ -122,4 +122,3 @@ python yolov8_cls.py  # Classification
 ## More Information
 
 See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
-
diff --git a/yolov8/include/block.h b/yolov8/include/block.h
@@ -1,21 +1,29 @@
 #pragma once
 #include <map>
-#include <vector>
 #include <string>
+#include <vector>
 #include "NvInfer.h"
 
 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 ch, int k, int s, int p, std::string lname);
+nvinfer1::IElementWiseLayer* convBnSiLU(nvinfer1::INetworkDefinition* network,
+                                        std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input,
+                                        int ch, int k, int s, int p, std::string lname);
+
+nvinfer1::IElementWiseLayer* C2F(nvinfer1::INetworkDefinition* network,
+                                 std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int c1,
+                                 int c2, int n, bool shortcut, float e, std::string lname);
 
-nvinfer1::IElementWiseLayer* C2F(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap, 
-nvinfer1::ITensor& input, int c1, int c2, int n, bool shortcut, float e, std::string lname);
+nvinfer1::IElementWiseLayer* C2(nvinfer1::INetworkDefinition* network,
+                                std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c1,
+                                int c2, int n, bool shortcut, float e, std::string lname);
 
-nvinfer1::IElementWiseLayer* SPPF(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap, 
-nvinfer1::ITensor& input, int c1, int c2, int k, std::string lname);
+nvinfer1::IElementWiseLayer* SPPF(nvinfer1::INetworkDefinition* network,
+                                  std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int c1,
+                                  int c2, int k, std::string lname);
 
-nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap, 
-nvinfer1::ITensor& input, int ch, int grid, int k, int s, int p, std::string lname);
+nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
+                             nvinfer1::ITensor& input, int ch, int grid, int k, int s, int p, std::string lname);
 
-nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation = false);
+nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition* network,
+                                       std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation = false);
diff --git a/yolov8/include/model.h b/yolov8/include/model.h
@@ -1,13 +1,19 @@
 #pragma once
-#include "NvInfer.h"
-#include <string>
 #include <assert.h>
+#include <string>
+#include "NvInfer.h"
+
+nvinfer1::IHostMemory* buildEngineYolov8Det(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config,
+                                            nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw,
+                                            int& max_channels);
 
-nvinfer1::IHostMemory* buildEngineYolov8Det(nvinfer1::IBuilder* builder,
-nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw, int& max_channels);
+nvinfer1::IHostMemory* buildEngineYolov8DetP6(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config,
+                                              nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw,
+                                              int& max_channels);
 
-nvinfer1::IHostMemory* buildEngineYolov8Cls(nvinfer1::IBuilder* builder,
-nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw);
+nvinfer1::IHostMemory* buildEngineYolov8Cls(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config,
+                                            nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw);
 
-nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder,
-nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw, int& max_channels);
+nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config,
+                                            nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw,
+                                            int& max_channels);