Refactor yolov7 (wang-xinyu#1155)

* separate model and block

* move thresh to config

* use trt8 api

* move batchsize to config

* update readme

* refactor batch infer logic

* move max input image size to config

Author: wang-xinyu

yolov7/README.md

@@ -1,26 +1,30 @@
-# yolov7
+# YOLOv7
 
 The Pytorch implementation is [WongKinYiu/yolov7](https://github.com/WongKinYiu/yolov7).
 
 The tensorrt code is derived from [QIANXUNZDL123/tensorrtx-yolov7](https://github.com/QIANXUNZDL123/tensorrtx-yolov7)
 
+## Contributors
+
+<a href="https://github.com/QIANXUNZDL123"><img src="https://avatars.githubusercontent.com/u/46549527?v=4?s=48" width="40px;" alt=""/></a>
+<a href="https://github.com/lindsayshuo"><img src="https://avatars.githubusercontent.com/u/45239466?v=4?s=48" width="40px;" alt=""/></a>
+<a href="https://github.com/wang-xinyu"><img src="https://avatars.githubusercontent.com/u/15235574?s=48&v=4" width="40px;" alt=""/></a> 
+
+## Requirements
+
+- TensorRT 8.0+
+- OpenCV 3.4.0+
+
 ## Different versions of yolov7
 
 Currently, we support yolov7 v0.1
 
 - For yolov7 v0.1, download .pt from [yolov7 release v0.1](https://github.com/WongKinYiu/yolov7/releases/tag/v0.10), then follow how-to-run in current page.
 
-
 ## Config
 
 - Choose the model tiny/v7/x/d6/w6/e6/e6e from command line arguments.
-- Input shape defined in yololayer.h
-- Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model**
-- INT8/FP16/FP32 can be selected by the macro in yolov7.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below**
-- GPU id can be selected by the macro in yolov7.cpp
-- NMS thresh in yolov7.cpp
-- BBox confidence thresh in yolov7.cpp
-- Batch size in yolov7.cpp
+- Check more configs in [include/config.h](./include/config.h)
 
 ## How to Run, yolov7-tiny as example
 
@@ -38,7 +42,7 @@ python gen_wts.py
 
 ```
 cd {tensorrtx}/yolov7/
-// update CLASS_NUM in yololayer.h if your model is trained on custom dataset
+// update kNumClass in config.h if your model is trained on custom dataset
 mkdir build
 cd build
 cp {WongKinYiu}/yolov7/yolov7.wts {tensorrtx}/yolov7/build
@@ -67,18 +71,15 @@ python yolov7_trt.py
 
 2. unzip it in yolov7/build
 
-3. set the macro `USE_INT8` in yolov7.cpp and make
+3. set the macro `USE_INT8` in config.h and make
 
 4. serialize the model and test
 
 <p align="center">
-<img src="https://user-images.githubusercontent.com/15235574/78247927-4d9fac00-751e-11ea-8b1b-704a0aeb3fcf.jpg">
-</p>
-
-<p align="center">
-<img src="https://user-images.githubusercontent.com/15235574/78247970-60b27c00-751e-11ea-88df-41473fed4823.jpg">
+<img src="https://user-images.githubusercontent.com/15235574/78247927-4d9fac00-751e-11ea-8b1b-704a0aeb3fcf.jpg" height="360px;">
 </p>
 
 ## More Information
 
 See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
+

yolov7/include/config.h

@@ -1,5 +1,12 @@
 #pragma once
 
+/* --------------------------------------------------------
+ * These configs are related to tensorrt model, if these are changed,
+ * please re-compile and re-serialize the tensorrt model.
+ * --------------------------------------------------------*/
+
+// For INT8, you need prepare the calibration dataset, please refer to
+// https://github.com/wang-xinyu/tensorrtx/tree/master/yolov7#int8-quantization
 #define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
 
 // These are used to define input/output tensor names,
@@ -8,6 +15,7 @@ const static char* kInputTensorName = "data";
 const static char* kOutputTensorName = "prob";
 
 const static int kNumClass = 80;
+const static int kBatchSize = 1;
 
 // Yolo's input width and height must by divisible by 32
 const static int kInputH = 640;
@@ -19,6 +27,20 @@ const static int kMaxNumOutputBbox = 1000;
 
 const static int kNumAnchor = 3;
 
-// The bboxes whose conf lower kIgnoreThresh will be ignored in yololayer plugin.
+// The bboxes whose confidence is lower than kIgnoreThresh will be ignored in yololayer plugin.
 const static float kIgnoreThresh = 0.1f;
 
+/* --------------------------------------------------------
+ * These configs are not related to tensorrt model, if these are changed,
+ * please re-compile, but no need to re-serialize the tensorrt model.
+ * --------------------------------------------------------*/
+
+// NMS overlapping thresh and final detection confidence thresh
+const static float kNmsThresh = 0.45f;
+const static float kConfThresh = 0.5f;
+
+const static int kGpuId = 0;
+
+// If your image size is larger than 4096 * 3112, please increase this value
+const static int kMaxInputImageSize = 4096 * 3112;
+

yolov7/include/model.h

@@ -3,10 +3,10 @@
 #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);
+nvinfer1::IHostMemory* build_engine_yolov7e6e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::IHostMemory* build_engine_yolov7d6(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::IHostMemory* build_engine_yolov7e6(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::IHostMemory* build_engine_yolov7w6(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::IHostMemory* build_engine_yolov7x(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::IHostMemory* build_engine_yolov7(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
+nvinfer1::IHostMemory* build_engine_yolov7_tiny(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name);

yolov7/include/postprocess.h

@@ -7,3 +7,7 @@ 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);
 
+void batch_nms(std::vector<std::vector<Detection>>& batch_res, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh = 0.5);
+
+void draw_bbox(std::vector<cv::Mat>& img_batch, std::vector<std::vector<Detection>>& res_batch);
+

yolov7/include/preprocess.h

@@ -1,16 +1,15 @@
-#ifndef __PREPROCESS_H
-#define __PREPROCESS_H
+#pragma once
 
 #include <cuda_runtime.h>
 #include <cstdint>
-
-
-struct AffineMatrix{
-    float value[6];
-};
-
-
-void preprocess_kernel_img(uint8_t* src, int src_width, int src_height,
+#include <opencv2/opencv.hpp>
+
+void cuda_preprocess_init(int max_image_size);
+void cuda_preprocess_destroy();
+void cuda_preprocess(uint8_t* src, int src_width, int src_height,
+                     float* dst, int dst_width, int dst_height,
+                     cudaStream_t stream);
+void cuda_batch_preprocess(std::vector<cv::Mat>& img_batch,
                            float* dst, int dst_width, int dst_height,
                            cudaStream_t stream);
-#endif  // __PREPROCESS_H
+