set default FP16 and update GELAN (#1494)

* set default FP16 and update GELAN * use pre-commit * use pre-commit-V2
2024-04-23 18:39:10 +08:00 · 2024-04-23 18:39:10 +08:00 · bfe13faac7
commit bfe13faac7
parent d033a63f07
10 changed files with 444 additions and 63 deletions
--- a/yolov9/README.md
+++ b/yolov9/README.md
@ -6,15 +6,24 @@ The Pytorch implementation is [WongKinYiu/yolov9](https://github.com/WongKinYiu/

 <a href="https://github.com/WuxinrongY"><img src="https://avatars.githubusercontent.com/u/53141838?v=4?s=48" width="40px;" alt=""/></a>

+
 ## Progress
 - [x] YOLOv9-c:
-    - [x] FP32
-    - [x] FP16
-    - [x] INT8
+  - [x] FP32
+  - [x] FP16
+  - [x] INT8
 - [x] YOLOv9-e:
-    - [x] FP32
-    - [x] FP16
-    - [x] INT8
+  - [x] FP32
+  - [x] FP16
+  - [x] INT8
+- [x] GELAN-c:
+  - [x] FP32
+  - [x] FP16
+  - [x] INT8
+- [x] GELAN-e:
+  - [x] FP32
+  - [x] FP16
+  - [x] INT8

 ## Requirements

@ -32,7 +41,10 @@ The speed test is done on a desktop with R7-5700G CPU and RTX 4060Ti GPU. The in
 | tensorrt | YOLOv9-c | 13.5ms | 4.6ms | 3.0ms |
 | tensorrt | YOLOv9-e | 8.3ms | 3.2ms | 2.15ms |

+**GELAN will be updated later.**
+
 YOLOv9-e is faster than YOLOv9-c in tensorrt, because the YOLOv9-e requires fewer layers of inference.
+
 ```
 YOLOv9-c:
 [[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]] # [A3, A4, A5, P3, P4, P5]
@ -42,7 +54,7 @@ YOLOv9-e:

 ```

-In DualDDetect, the A3, A4, A5, P3, P4, P5 are the output of the backbone. The first 3 layers are used for the inference of the final result. 
+In DualDDetect, the A3, A4, A5, P3, P4, P5 are the output of the backbone. The first 3 layers are used for the inference of the final result.

 The YOLOv9-c requires 37 layers of inference, but YOLOv9-e requires 35 layers of inference.

@ -103,5 +115,3 @@ python yolov9_trt.py
 ## More Information

 See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
-
-
--- a/yolov9/demo.cpp
+++ b/yolov9/demo.cpp
@ -9,10 +9,8 @@
 #include "utils.h"

 using namespace nvinfer1;
-
 const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1;
 static Logger gLogger;
-
 void serialize_engine(unsigned int max_batchsize, std::string& wts_name, std::string& sub_type,
                      std::string& engine_name) {
    // Create builder
@ -25,6 +23,10 @@ void serialize_engine(unsigned int max_batchsize, std::string& wts_name, std::st
        serialized_engine = build_engine_yolov9_e(max_batchsize, builder, config, DataType::kFLOAT, wts_name);
    } else if (sub_type == "c") {
        serialized_engine = build_engine_yolov9_c(max_batchsize, builder, config, DataType::kFLOAT, wts_name);
+    } else if (sub_type == "ge") {
+        serialized_engine = build_engine_gelan_e(max_batchsize, builder, config, DataType::kFLOAT, wts_name);
+    } else if (sub_type == "gc") {
+        serialized_engine = build_engine_gelan_c(max_batchsize, builder, config, DataType::kFLOAT, wts_name);
    } else {
        return;
    }
@ -113,7 +115,7 @@ int main(int argc, char** argv) {

    std::string wts_name = "";
    std::string engine_name = "";
-    std::string img_dir;
+    std::string img_dir = "";
    std::string sub_type = "";
    // speed test or inference
    // const int speed_test_iter = 1000;
@ -121,7 +123,7 @@ int main(int argc, char** argv) {

    if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type)) {
        std::cerr << "Arguments not right!" << std::endl;
-        std::cerr << "./yolov9 -s [.wts] [.engine] [c/e]  // serialize model to plan file" << std::endl;
+        std::cerr << "./yolov9 -s [.wts] [.engine] [c/e/gc/ge]  // serialize model to plan file" << std::endl;
        std::cerr << "./yolov9 -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
        return -1;
    }
--- a/yolov9/include/block.h
+++ b/yolov9/include/block.h
@ -42,3 +42,5 @@ nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std
 nvinfer1::ILayer* convBnNoAct(nvinfer1::INetworkDefinition* network,
                              std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int ch,
                              int k, int s, int p, std::string lname, int g);
+std::vector<IConcatenationLayer*> DDetect(INetworkDefinition* network, std::map<std::string, Weights>& weightMap,
+                                          std::vector<ILayer*> dets, int cls, std::vector<int> ch, std::string lname);
--- a/yolov9/include/config.h
+++ b/yolov9/include/config.h
@ -7,8 +7,7 @@

 // For INT8, you need prepare the calibration dataset, please refer to
 // https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5#int8-quantization
-#define USE_INT8  // set USE_INT8 or USE_FP16 or USE_FP32
-
+#define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
 #ifdef USE_INT8
 const static char* gCalibTablePath = "./calib";
 #endif
@ -56,4 +55,3 @@ const static int kGpuId = 0;

 // If your image size is larger than 4096 * 3112, please increase this value
 const static int kMaxInputImageSize = 4096 * 3112;
-
--- a/yolov9/include/model.h
+++ b/yolov9/include/model.h
@ -2,5 +2,16 @@

 #include <NvInfer.h>
 #include <string>
-nvinfer1::IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name);
-nvinfer1::IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name);
+nvinfer1::IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
+                                             nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
+                                             std::string& wts_name);
+nvinfer1::IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
+                                             nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
+                                             std::string& wts_name);
+
+nvinfer1::IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
+                                            nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
+                                            std::string& wts_name);
+nvinfer1::IHostMemory* build_engine_gelan_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
+                                            nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
+                                            std::string& wts_name);
--- a/yolov9/include/utils.h
+++ b/yolov9/include/utils.h
@ -1,23 +1,22 @@
 #pragma once

 #include <dirent.h>
-#include <fstream>
-#include <unordered_map>
-#include <string>
-#include <sstream>
-#include <vector>
 #include <cstring>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+#include <vector>

 static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::string>& file_names) {
-    DIR *p_dir = opendir(p_dir_name);
+    DIR* p_dir = opendir(p_dir_name);
    if (p_dir == nullptr) {
        return -1;
    }

    struct dirent* p_file = nullptr;
    while ((p_file = readdir(p_dir)) != nullptr) {
-        if (strcmp(p_file->d_name, ".") != 0 &&
-            strcmp(p_file->d_name, "..") != 0) {
+        if (strcmp(p_file->d_name, ".") != 0 && strcmp(p_file->d_name, "..") != 0) {
            //std::string cur_file_name(p_dir_name);
            //cur_file_name += "/";
            //cur_file_name += p_file->d_name;
@ -29,7 +28,6 @@ static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::str
    closedir(p_dir);
    return 0;
 }
-
 // Function to trim leading and trailing whitespace from a string
 static inline std::string trim_leading_whitespace(const std::string& str) {
    size_t first = str.find_first_not_of(' ');
@ -67,4 +65,3 @@ static inline int read_labels(const std::string labels_filename, std::unordered_

    return 0;
 }
-
--- a/yolov9/src/block.cpp
+++ b/yolov9/src/block.cpp
@ -14,7 +14,6 @@
 #include <vector>

 using namespace nvinfer1;
-
 // TensorRT weight files have a simple space delimited format:
 // [type] [size] <data x size in hex>
 void PrintDim(const ILayer* layer, std::string log) {
@ -423,3 +422,35 @@ std::vector<IConcatenationLayer*> DualDDetect(INetworkDefinition* network, std::
    }
    return ret;
 }
+
+std::vector<IConcatenationLayer*> DDetect(INetworkDefinition* network, std::map<std::string, Weights>& weightMap,
+                                          std::vector<ILayer*> dets, int cls, std::vector<int> ch, std::string lname) {
+    int c2 = std::max(int(ch[0] / 4), int(16 * 4));
+    int c3 = std::max(ch[0], std::min(cls * 2, 128));
+    int reg_max = 16;
+
+    std::vector<ILayer*> bboxlayers;
+    std::vector<ILayer*> clslayers;
+
+    for (int i = 0; i < dets.size(); i++) {
+        // Conv(x, c2, 3), Conv(c2, c2, 3, g=4), nn.Conv2d(c2, 4 * self.reg_max, 1, groups=4)
+        bboxlayers.push_back(DetectBbox_Conv(network, weightMap, *dets[i]->getOutput(0), c2, reg_max,
+                                             lname + ".cv2." + std::to_string(i)));
+        // Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, self.nc, 1)
+        auto cls_layer = DetectCls_Conv(network, weightMap, *dets[i]->getOutput(0), c3, cls,
+                                        lname + ".cv3." + std::to_string(i));
+        auto dim = cls_layer->getOutput(0)->getDimensions();
+        nvinfer1::IShuffleLayer* shuffle = network->addShuffle(*cls_layer->getOutput(0));
+        shuffle->setReshapeDimensions(nvinfer1::Dims2{kNumClass, dim.d[1] * dim.d[2]});
+        clslayers.push_back(shuffle);
+    }
+
+    std::vector<IConcatenationLayer*> ret;
+    for (int i = 0; i < dets.size(); i++) {
+        // softmax 16*4, w, h => 16, 4, w, h
+        auto loc = DFL(network, weightMap, *bboxlayers[i]->getOutput(0), 16, 1, 1, 0, lname + ".dfl");
+        nvinfer1::ITensor* inputTensor[] = {loc->getOutput(0), clslayers[i]->getOutput(0)};
+        ret.push_back(network->addConcatenation(inputTensor, 2));
+    }
+    return ret;
+}
--- a/yolov9/src/calibrator.cpp
+++ b/yolov9/src/calibrator.cpp
@ -2,12 +2,11 @@
 #include "cuda_utils.h"
 #include "utils.h"

+#include <fstream>
 #include <iostream>
 #include <iterator>
-#include <fstream>
-#include <opencv2/opencv.hpp>
 #include <opencv2/dnn/dnn.hpp>
-
+#include <opencv2/opencv.hpp>
 static cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
    int w, h, x, y;
    float r_w = input_w / (img.cols * 1.0);
@ -30,7 +29,9 @@ static cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
    return out;
 }

-Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache)
+Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir,
+                                               const char* calib_table_name, const char* input_blob_name,
+                                               bool read_cache)
    : batchsize_(batchsize),
      input_w_(input_w),
      input_h_(input_h),
@ -69,7 +70,8 @@ bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int
        input_imgs_.push_back(pr_img);
    }
    img_idx_ += batchsize_;
-    cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false);
+    cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0),
+                                           true, false);

    CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr<float>(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice));
    assert(!strcmp(names[0], input_blob_name_));
@ -94,4 +96,3 @@ void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t len
    std::ofstream output(calib_table_name_, std::ios::binary);
    output.write(reinterpret_cast<const char*>(cache), length);
 }
-
--- a/yolov9/src/model.cpp
+++ b/yolov9/src/model.cpp
@ -11,7 +11,6 @@
 #include "yololayer.h"

 using namespace nvinfer1;
-
 #ifdef USE_INT8
 void Calibrator(IBuilder* builder, IBuilderConfig* config) {
    std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
@ -432,3 +431,325 @@ IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, IBuilder* builder,

    return serialized_model;
 }
+
+nvinfer1::IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
+                                            nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
+                                            std::string& wts_name) {
+    /* ------ Create the builder ------ */
+    INetworkDefinition* network = builder->createNetworkV2(0U);
+
+    ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW});
+    assert(data);
+    std::map<std::string, Weights> weightMap = loadWeights(wts_name);
+
+    /* ------backbone------ */
+    // [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
+    auto conv_1 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.1", 1);
+    assert(conv_1);
+    // [-1, 1, Conv, [128, 3, 2]],  # 2-P2/4
+    auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 128, 3, 2, 1, "model.2");
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]],  # 3
+    auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 256, 128, 64, 2, "model.3");
+    // avg-conv down
+    // [-1, 1, ADown, [256]],  # 4-P3/8
+    auto adown_4 = ADown(network, weightMap, *repncspelan_3->getOutput(0), 256, "model.4");
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]],  # 5
+    auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 512, 256, 128, 2, "model.5");
+    // avg-conv down
+    // [-1, 1, ADown, [512]],  # 6-P4/16
+    auto adown_6 = ADown(network, weightMap, *repncspelan_5->getOutput(0), 512, "model.6");
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]],  # 7
+    auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 1024, 512, 256, 2, "model.7");
+    // avg-conv down
+    // [-1, 1, ADown, [1024]],  # 8-P5/32
+    auto adown_8 = ADown(network, weightMap, *repncspelan_7->getOutput(0), 1024, "model.8");
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]],  # 9
+    auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 1024, 512, 256, 2, "model.9");
+
+    // [1, 1, CBLinear, [[64]]], # 10
+    auto cblinear_10 = CBLinear(network, weightMap, *conv_1->getOutput(0), {64}, 1, 1, 0, 1, "model.10");
+    // [3, 1, CBLinear, [[64, 128]]], # 11
+    auto cblinear_11 = CBLinear(network, weightMap, *repncspelan_3->getOutput(0), {64, 128}, 1, 1, 0, 1, "model.11");
+    // [5, 1, CBLinear, [[64, 128, 256]]], # 12
+    auto cblinear_12 =
+            CBLinear(network, weightMap, *repncspelan_5->getOutput(0), {64, 128, 256}, 1, 1, 0, 1, "model.12");
+    // [7, 1, CBLinear, [[64, 128, 256, 512]]], # 13
+    auto cblinear_13 =
+            CBLinear(network, weightMap, *repncspelan_7->getOutput(0), {64, 128, 256, 512}, 1, 1, 0, 1, "model.13");
+    // [9, 1, CBLinear, [[64, 128, 256, 512, 1024]]], # 14
+    auto cblinear_14 = CBLinear(network, weightMap, *repncspelan_9->getOutput(0), {64, 128, 256, 512, 1024}, 1, 1, 0, 1,
+                                "model.14");
+
+    // conv down
+    // [0, 1, Conv, [64, 3, 2]],  # 15-P1/2
+    auto conv_15 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.15", 1);
+    // [[10, 11, 12, 13, 14, -1], 1, CBFuse, [[0, 0, 0, 0, 0]]], # 16
+    auto cbfuse_16 = CBFuse(
+            network, {cblinear_10, cblinear_11, cblinear_12, cblinear_13, cblinear_14, std::vector<ILayer*>{conv_15}},
+            {0, 0, 0, 0, 0, 0}, {2, 4, 8, 16, 32, 2});
+
+    // conv down
+    // [-1, 1, Conv, [128, 3, 2]],  # 17-P2/4
+    auto conv_17 = convBnSiLU(network, weightMap, *cbfuse_16->getOutput(0), 128, 3, 2, 1, "model.17");
+    // [[11, 12, 13, 14, -1], 1, CBFuse, [[1, 1, 1, 1]]], # 18
+    auto cbfuse_18 =
+            CBFuse(network, {cblinear_11, cblinear_12, cblinear_13, cblinear_14, std::vector<ILayer*>{conv_17}},
+                   {1, 1, 1, 1, 0}, {4, 8, 16, 32, 4});
+
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]],  # 19
+    auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cbfuse_18->getOutput(0), 128, 256, 128, 64, 2, "model.19");
+
+    // avg-conv down fuse
+    // [-1, 1, ADown, [256]],  # 20-P3/8
+    auto adown_20 = ADown(network, weightMap, *repncspelan_19->getOutput(0), 256, "model.20");
+    // [[12, 13, 14, -1], 1, CBFuse, [[2, 2, 2]]], # 21
+    auto cbfuse_21 = CBFuse(network, {cblinear_12, cblinear_13, cblinear_14, std::vector<ILayer*>{adown_20}},
+                            {2, 2, 2, 0}, {8, 16, 32, 8});
+
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]],  # 22
+    auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cbfuse_21->getOutput(0), 256, 512, 256, 128, 2, "model.22");
+
+    // avg-conv down fuse
+    // [-1, 1, ADown, [512]],  # 23-P4/16
+    auto adown_23 = ADown(network, weightMap, *repncspelan_22->getOutput(0), 512, "model.23");
+    // [[13, 14, -1], 1, CBFuse, [[3, 3]]], # 24
+    auto cbfuse_24 =
+            CBFuse(network, {cblinear_13, cblinear_14, std::vector<ILayer*>{adown_23}}, {3, 3, 0}, {16, 32, 16});
+
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]],  # 25
+    auto repncspelan_25 =
+            RepNCSPELAN4(network, weightMap, *cbfuse_24->getOutput(0), 512, 1024, 512, 256, 2, "model.25");
+
+    // avg-conv down fuse
+    // [-1, 1, ADown, [1024]],  # 26-P5/32
+    auto adown_26 = ADown(network, weightMap, *repncspelan_25->getOutput(0), 1024, "model.26");
+    // [[14, -1], 1, CBFuse, [[4]]], # 27
+    auto cbfuse_27 = CBFuse(network, {cblinear_14, std::vector<ILayer*>{adown_26}}, {4, 0}, {32, 32});
+
+    // csp-elan block
+    // [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]],  # 28
+    auto repncspelan_28 =
+            RepNCSPELAN4(network, weightMap, *cbfuse_27->getOutput(0), 512, 1024, 512, 256, 2, "model.28");
+
+    // elan-spp block
+    // [28, 1, SPPELAN, [512, 256]],  # 29
+    auto sppelan_29 = SPPELAN(network, weightMap, *repncspelan_28->getOutput(0), 1024, 512, 256, "model.29");
+
+    // # up-concat merge
+    // [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+    auto upsample_30 = network->addResize(*sppelan_29->getOutput(0));
+    upsample_30->setResizeMode(ResizeMode::kNEAREST);
+    const float scales_30[] = {1.0, 2.0, 2.0};
+    upsample_30->setScales(scales_30, 3);
+    // [[-1, 25], 1, Concat, [1]],  # cat backbone P4
+    ITensor* input_tensor_31[] = {upsample_30->getOutput(0), repncspelan_25->getOutput(0)};
+    auto cat_31 = network->addConcatenation(input_tensor_31, 2);
+
+    // # csp-elan block
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]],  # 32
+    auto repncspelan_32 = RepNCSPELAN4(network, weightMap, *cat_31->getOutput(0), 1536, 512, 512, 256, 2, "model.32");
+
+    // # up-concat merge
+    // [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+    auto upsample_33 = network->addResize(*repncspelan_32->getOutput(0));
+    upsample_33->setResizeMode(ResizeMode::kNEAREST);
+    const float scales_33[] = {1.0, 2.0, 2.0};
+    upsample_33->setScales(scales_33, 3);
+    // [[-1, 22], 1, Concat, [1]],  # cat backbone P3
+    ITensor* input_tensor_34[] = {upsample_33->getOutput(0), repncspelan_22->getOutput(0)};
+    auto cat_34 = network->addConcatenation(input_tensor_34, 2);
+
+    // # csp-elan block
+    // [-1, 1, RepNCSPELAN4, [256, 256, 128, 2]],  # 35
+    auto repncspelan_35 = RepNCSPELAN4(network, weightMap, *cat_34->getOutput(0), 1024, 256, 256, 128, 2, "model.35");
+
+    // # avg-conv-down merge
+    // [-1, 1, ADown, [256]],
+    auto adown_36 = ADown(network, weightMap, *repncspelan_35->getOutput(0), 256, "model.36");
+    // [[-1, 32], 1, Concat, [1]],  # cat head P4
+    ITensor* input_tensor_37[] = {adown_36->getOutput(0), repncspelan_32->getOutput(0)};
+    auto cat_37 = network->addConcatenation(input_tensor_37, 2);
+
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]],  # 38 (P4/16-medium)
+    auto repncspelan_38 = RepNCSPELAN4(network, weightMap, *cat_37->getOutput(0), 768, 512, 512, 256, 2, "model.38");
+
+    // # avg-conv-down merge
+    // [-1, 1, ADown, [512]],
+    auto adown_39 = ADown(network, weightMap, *repncspelan_38->getOutput(0), 512, "model.39");
+    // [[-1, 29], 1, Concat, [1]],  # cat head P5
+    ITensor* input_tensor_40[] = {adown_39->getOutput(0), sppelan_29->getOutput(0)};
+    auto cat_40 = network->addConcatenation(input_tensor_40, 2);
+
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 1024, 512, 2]],  # 41 (P5/32-large)
+    auto repncspelan_41 = RepNCSPELAN4(network, weightMap, *cat_40->getOutput(0), 1024, 512, 1024, 512, 2, "model.41");
+
+    auto ddetect_42 = DDetect(network, weightMap, std::vector<ILayer*>{repncspelan_35, repncspelan_38, repncspelan_41},
+                              kNumClass, {256, 512, 512}, "model.42");
+
+    nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, ddetect_42, false);
+    yolo->getOutput(0)->setName(kOutputTensorName);
+    network->markOutput(*yolo->getOutput(0));
+
+    builder->setMaxBatchSize(kBatchSize);
+    config->setMaxWorkspaceSize(16 * (1 << 20));
+
+#if defined(USE_FP16)
+    config->setFlag(nvinfer1::BuilderFlag::kFP16);
+#elif defined(USE_INT8)
+    std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
+    assert(builder->platformHasFastInt8());
+    config->setFlag(nvinfer1::BuilderFlag::kINT8);
+    auto* calibrator =
+            new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName);
+    config->setInt8Calibrator(calibrator);
+#endif
+
+    std::cout << "Building engine, please wait for a while..." << std::endl;
+    IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config);
+    std::cout << "Build engine successfully!" << std::endl;
+
+    delete network;
+
+    // Release host memory
+    for (auto& mem : weightMap) {
+        free((void*)(mem.second.values));
+    }
+
+    return serialized_model;
+}
+nvinfer1::IHostMemory* build_engine_gelan_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder,
+                                            nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
+                                            std::string& wts_name) {
+    /* ------ Create the builder ------ */
+    INetworkDefinition* network = builder->createNetworkV2(0U);
+
+    ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW});
+    assert(data);
+    std::map<std::string, Weights> weightMap = loadWeights(wts_name);
+
+    // # conv down
+    // [-1, 1, Conv, [64, 3, 2]],  # 1-P1/2
+    auto conv_1 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.0", 1);
+    // # conv down
+    // [-1, 1, Conv, [128, 3, 2]],  # 2-P2/4
+    auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 128, 3, 2, 1, "model.1");
+    // # elan-1 block
+    // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]],  # 3
+    auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 256, 128, 64, 1, "model.2");
+    // # avg-conv down
+    // [-1, 1, ADown, [256]],  # 4-P3/8
+    auto adown_4 = ADown(network, weightMap, *repncspelan_3->getOutput(0), 256, "model.3");
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]],  # 5
+    auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 512, 256, 128, 1, "model.4");
+    // # avg-conv down
+    // [-1, 1, ADown, [512]],  # 6-P4/16
+    auto adown_6 = ADown(network, weightMap, *repncspelan_5->getOutput(0), 512, "model.5");
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 7
+    auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 512, 512, 256, 1, "model.6");
+    // # avg-conv down
+    // [-1, 1, ADown, [512]],  # 8-P5/32
+    auto adown_8 = ADown(network, weightMap, *repncspelan_7->getOutput(0), 512, "model.7");
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 9
+    auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 512, 512, 256, 1, "model.8");
+    // # elan-spp block
+    // [-1, 1, SPPELAN, [512, 256]],  # 10
+    auto sppelan_10 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 512, 256, "model.9");
+
+    // # up-concat merge
+    // [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+    auto upsample_11 = network->addResize(*sppelan_10->getOutput(0));
+    upsample_11->setResizeMode(ResizeMode::kNEAREST);
+    const float scales_11[] = {1.0, 2.0, 2.0};
+    upsample_11->setScales(scales_11, 3);
+    // [[-1, 7], 1, Concat, [1]],  # cat backbone P4
+    ITensor* input_tensor_12[] = {upsample_11->getOutput(0), repncspelan_7->getOutput(0)};
+    auto cat_12 = network->addConcatenation(input_tensor_12, 2);
+
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 13
+    auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 1536, 512, 512, 256, 1, "model.12");
+
+    // # up-concat merge
+    // [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+    auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0));
+    upsample_14->setResizeMode(ResizeMode::kNEAREST);
+    const float scales_14[] = {1.0, 2.0, 2.0};
+    upsample_14->setScales(scales_14, 3);
+    // [[-1, 5], 1, Concat, [1]],  # cat backbone P3
+    ITensor* input_tensor_15[] = {upsample_14->getOutput(0), repncspelan_5->getOutput(0)};
+    auto cat_15 = network->addConcatenation(input_tensor_15, 2);
+
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]],  # 16 (P3/8-small)
+    auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 1024, 256, 256, 128, 1, "model.15");
+
+    // # avg-conv-down merge
+    // [-1, 1, ADown, [256]],
+    auto adown_17 = ADown(network, weightMap, *repncspelan_16->getOutput(0), 256, "model.16");
+    // [[-1, 13], 1, Concat, [1]],  # cat head P4
+    ITensor* input_tensor_18[] = {adown_17->getOutput(0), repncspelan_13->getOutput(0)};
+    auto cat_18 = network->addConcatenation(input_tensor_18, 2);
+
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 19 (P4/16-medium)
+    auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 768, 512, 512, 256, 1, "model.18");
+
+    // # avg-conv-down merge
+    // [-1, 1, ADown, [512]],
+    auto adown_20 = ADown(network, weightMap, *repncspelan_19->getOutput(0), 512, "model.19");
+    // [[-1, 10], 1, Concat, [1]],  # cat head P5
+    ITensor* input_tensor_21[] = {adown_20->getOutput(0), sppelan_10->getOutput(0)};
+    auto cat_21 = network->addConcatenation(input_tensor_21, 2);
+
+    // # elan-2 block
+    // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]],  # 22 (P5/32-large)
+    auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 1024, 512, 512, 256, 1, "model.21");
+
+    // # detection head
+    // # detect
+    // [[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]],  # DualDDetect(A3, A4, A5, P3, P4, P5)
+    auto ddetect_23 = DDetect(network, weightMap, std::vector<ILayer*>{repncspelan_16, repncspelan_19, repncspelan_22},
+                              kNumClass, {256, 512, 512}, "model.22");
+
+    nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, ddetect_23, false);
+    yolo->getOutput(0)->setName(kOutputTensorName);
+    network->markOutput(*yolo->getOutput(0));
+
+    builder->setMaxBatchSize(kBatchSize);
+    config->setMaxWorkspaceSize(16 * (1 << 20));
+
+#if defined(USE_FP16)
+    config->setFlag(nvinfer1::BuilderFlag::kFP16);
+#elif defined(USE_INT8)
+    std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
+    assert(builder->platformHasFastInt8());
+    config->setFlag(nvinfer1::BuilderFlag::kINT8);
+    auto* calibrator =
+            new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName);
+    config->setInt8Calibrator(calibrator);
+#endif
+
+    std::cout << "Building engine, please wait for a while..." << std::endl;
+    IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config);
+    std::cout << "Build engine successfully!" << std::endl;
+
+    delete network;
+
+    // Release host memory
+    for (auto& mem : weightMap) {
+        free((void*)(mem.second.values));
+    }
+
+    return serialized_model;
+}
--- a/yolov9/src/postprocess.cpp
+++ b/yolov9/src/postprocess.cpp
@ -1,6 +1,5 @@
 #include "postprocess.h"
 #include "utils.h"
-
 cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
    float l, r, t, b;
    float r_w = kInputW / (img.cols * 1.0);
@ -29,16 +28,16 @@ cv::Rect get_rect(cv::Mat& img, float bbox[4]) {

 static float iou(float lbox[4], float rbox[4]) {
    float interBox[] = {
-            (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left
-            (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right
-            (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top
-            (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom
+            (std::max)(lbox[0] - lbox[2] / 2.f, rbox[0] - rbox[2] / 2.f),  //left
+            (std::min)(lbox[0] + lbox[2] / 2.f, rbox[0] + rbox[2] / 2.f),  //right
+            (std::max)(lbox[1] - lbox[3] / 2.f, rbox[1] - rbox[3] / 2.f),  //top
+            (std::min)(lbox[1] + lbox[3] / 2.f, rbox[1] + rbox[3] / 2.f),  //bottom
    };

    if (interBox[2] > interBox[3] || interBox[0] > interBox[1])
        return 0.0f;

-    float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]);
+    float interBoxS = (interBox[1] - interBox[0]) * (interBox[3] - interBox[2]);
    return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
 }

@ -50,15 +49,17 @@ void nms(std::vector<Detection>& res, float* output, float conf_thresh, float nm
    int det_size = sizeof(Detection) / sizeof(float);
    std::map<float, std::vector<Detection>> m;
    for (int i = 0; i < output[0] && i < kMaxNumOutputBbox; i++) {
-        if (output[1 + det_size * i + 4] <= conf_thresh) continue;
+        if (output[1 + det_size * i + 4] <= conf_thresh)
+            continue;
        Detection det;
        memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float));
-        if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector<Detection>());
+        if (m.count(det.class_id) == 0)
+            m.emplace(det.class_id, std::vector<Detection>());
        // x1x2y1y2 -> xywh
-        float c_x = (det.bbox[0]+det.bbox[2])/2;
-        float c_y = (det.bbox[1]+det.bbox[3])/2;
-        float w = det.bbox[2]-det.bbox[0];
-        float h = det.bbox[3]-det.bbox[1];
+        float c_x = (det.bbox[0] + det.bbox[2]) / 2;
+        float c_y = (det.bbox[1] + det.bbox[3]) / 2;
+        float w = det.bbox[2] - det.bbox[0];
+        float h = det.bbox[3] - det.bbox[1];
        det.bbox[0] = c_x;
        det.bbox[1] = c_y;
        det.bbox[2] = w;
@ -81,7 +82,8 @@ void nms(std::vector<Detection>& res, float* output, float conf_thresh, float nm
    }
 }

-void batch_nms(std::vector<std::vector<Detection>>& res_batch, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh) {
+void batch_nms(std::vector<std::vector<Detection>>& res_batch, float* output, int batch_size, int output_size,
+               float conf_thresh, float nms_thresh) {
    res_batch.resize(batch_size);
    for (int i = 0; i < batch_size; i++) {
        nms(res_batch[i], &output[i * output_size], conf_thresh, nms_thresh);
@ -95,12 +97,14 @@ void draw_bbox(std::vector<cv::Mat>& img_batch, std::vector<std::vector<Detectio
        for (size_t j = 0; j < res.size(); j++) {
            cv::Rect r = get_rect(img, res[j].bbox);
            cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
-            cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
+            cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2,
+                        cv::Scalar(0xFF, 0xFF, 0xFF), 2);
        }
    }
    // draw num of objets to img
    for (size_t i = 0; i < img_batch.size(); i++) {
-        cv::putText(img_batch[i], std::to_string(res_batch[i].size()), cv::Point(0, 20), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
+        cv::putText(img_batch[i], std::to_string(res_batch[i].size()), cv::Point(0, 20), cv::FONT_HERSHEY_PLAIN, 1.2,
+                    cv::Scalar(0xFF, 0xFF, 0xFF), 2);
    }
 }

@ -158,11 +162,11 @@ cv::Mat scale_mask(cv::Mat mask, cv::Mat img) {
    return res;
 }

-void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map) {
-    static std::vector<uint32_t> colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A,
-                                           0x92CC17, 0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF,
-                                           0x344593, 0x6473FF, 0x0018EC, 0x8438FF, 0x520085, 0xCB38FF,
-                                           0xFF95C8, 0xFF37C7};
+void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::Mat>& masks,
+                    std::unordered_map<int, std::string>& labels_map) {
+    static std::vector<uint32_t> colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A, 0x92CC17,
+                                           0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF, 0x344593, 0x6473FF,
+                                           0x0018EC, 0x8438FF, 0x520085, 0xCB38FF, 0xFF95C8, 0xFF37C7};
    for (size_t i = 0; i < dets.size(); i++) {
        cv::Mat img_mask = scale_mask(masks[i], img);
        auto color = colors[(int)dets[i].class_id % colors.size()];
@ -172,7 +176,8 @@ void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::
        for (int x = r.x; x < r.x + r.width; x++) {
            for (int y = r.y; y < r.y + r.height; y++) {
                float val = img_mask.at<float>(y, x);
-                if (val <= 0.5) continue;
+                if (val <= 0.5)
+                    continue;
                img.at<cv::Vec3b>(y, x)[0] = img.at<cv::Vec3b>(y, x)[0] / 2 + bgr[0] / 2;
                img.at<cv::Vec3b>(y, x)[1] = img.at<cv::Vec3b>(y, x)[1] / 2 + bgr[1] / 2;
                img.at<cv::Vec3b>(y, x)[2] = img.at<cv::Vec3b>(y, x)[2] / 2 + bgr[2] / 2;
@ -182,7 +187,9 @@ void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::
        cv::rectangle(img, r, bgr, 2);

        // Get the size of the text
-        cv::Size textSize = cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL);
+        cv::Size textSize =
+                cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf),
+                                cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL);
        // Set the top left corner of the rectangle
        cv::Point topLeft(r.x, r.y - textSize.height);

@ -195,11 +202,12 @@ void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::
        // Draw the rectangle on the image
        cv::rectangle(img, topLeft, bottomRight, bgr, -1);

-        cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2);
-
+        cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf),
+                    cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2);
    }
 }
-void process_decode_ptr_host(std::vector<Detection> &res, const float* decode_ptr_host, int bbox_element, cv::Mat& img, int count) {
+void process_decode_ptr_host(std::vector<Detection>& res, const float* decode_ptr_host, int bbox_element, cv::Mat& img,
+                             int count) {
    Detection det;
    for (int i = 0; i < count; i++) {
        int basic_pos = 1 + i * bbox_element;
@ -215,7 +223,8 @@ void process_decode_ptr_host(std::vector<Detection> &res, const float* decode_pt
        }
    }
 }
-void batch_process(std::vector<std::vector<Detection>> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector<cv::Mat>& img_batch) {
+void batch_process(std::vector<std::vector<Detection>>& res_batch, const float* decode_ptr_host, int batch_size,
+                   int bbox_element, const std::vector<cv::Mat>& img_batch) {
    res_batch.resize(batch_size);
    int count = static_cast<int>(*decode_ptr_host);
    count = count > kMaxNumOutputBbox ? kMaxNumOutputBbox : count;
@ -225,4 +234,3 @@ void batch_process(std::vector<std::vector<Detection>> &res_batch, const float*
        process_decode_ptr_host(res_batch[i], &decode_ptr_host[i * count], bbox_element, img, count);
    }
 }
-