diff --git a/yolov9/README.md b/yolov9/README.md
index 37604be..90586b6 100644
--- a/yolov9/README.md
+++ b/yolov9/README.md
@@ -6,15 +6,24 @@ The Pytorch implementation is [WongKinYiu/yolov9](https://github.com/WongKinYiu/
+
## 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)
-
-
diff --git a/yolov9/demo.cpp b/yolov9/demo.cpp
index 2847acd..d193f8c 100644
--- 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;
}
diff --git a/yolov9/include/block.h b/yolov9/include/block.h
index 3b02d4b..ca7f79f 100644
--- a/yolov9/include/block.h
+++ b/yolov9/include/block.h
@@ -42,3 +42,5 @@ nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, int ch,
int k, int s, int p, std::string lname, int g);
+std::vector DDetect(INetworkDefinition* network, std::map& weightMap,
+ std::vector dets, int cls, std::vector ch, std::string lname);
diff --git a/yolov9/include/config.h b/yolov9/include/config.h
index 5863fef..2695a06 100644
--- 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;
-
diff --git a/yolov9/include/model.h b/yolov9/include/model.h
index eeb0e7c..ced2694 100644
--- a/yolov9/include/model.h
+++ b/yolov9/include/model.h
@@ -2,5 +2,16 @@
#include
#include
-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);
\ No newline at end of file
+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);
diff --git a/yolov9/include/utils.h b/yolov9/include/utils.h
index 2dea946..cd65de0 100644
--- a/yolov9/include/utils.h
+++ b/yolov9/include/utils.h
@@ -1,23 +1,22 @@
#pragma once
#include
-#include
-#include
-#include
-#include
-#include
#include
+#include
+#include
+#include
+#include
+#include
static inline int read_files_in_dir(const char* p_dir_name, std::vector& 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
using namespace nvinfer1;
-
// TensorRT weight files have a simple space delimited format:
// [type] [size]
void PrintDim(const ILayer* layer, std::string log) {
@@ -423,3 +422,35 @@ std::vector DualDDetect(INetworkDefinition* network, std::
}
return ret;
}
+
+std::vector DDetect(INetworkDefinition* network, std::map& weightMap,
+ std::vector dets, int cls, std::vector 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 bboxlayers;
+ std::vector 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 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;
+}
diff --git a/yolov9/src/calibrator.cpp b/yolov9/src/calibrator.cpp
index 7768d63..becafa5 100644
--- a/yolov9/src/calibrator.cpp
+++ b/yolov9/src/calibrator.cpp
@@ -2,12 +2,11 @@
#include "cuda_utils.h"
#include "utils.h"
+#include
#include
#include
-#include
-#include
#include
-
+#include
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(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(cache), length);
}
-
diff --git a/yolov9/src/model.cpp b/yolov9/src/model.cpp
index 18954b5..d450fb5 100644
--- 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 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{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{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{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{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{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{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 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{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;
+}
diff --git a/yolov9/src/postprocess.cpp b/yolov9/src/postprocess.cpp
index be25b62..9b2b7d2 100644
--- 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& res, float* output, float conf_thresh, float nm
int det_size = sizeof(Detection) / sizeof(float);
std::map> 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());
+ if (m.count(det.class_id) == 0)
+ m.emplace(det.class_id, std::vector());
// 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& res, float* output, float conf_thresh, float nm
}
}
-void batch_nms(std::vector>& res_batch, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh) {
+void batch_nms(std::vector>& 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& img_batch, std::vector& dets, std::vector& masks, std::unordered_map& labels_map) {
- static std::vector 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& dets, std::vector& masks,
+ std::unordered_map& labels_map) {
+ static std::vector 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& dets, std::vector(y, x);
- if (val <= 0.5) continue;
+ if (val <= 0.5)
+ continue;
img.at(y, x)[0] = img.at(y, x)[0] / 2 + bgr[0] / 2;
img.at(y, x)[1] = img.at(y, x)[1] / 2 + bgr[1] / 2;
img.at(y, x)[2] = img.at(y, x)[2] / 2 + bgr[2] / 2;
@@ -182,7 +187,9 @@ void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& dets, std::vector &res, const float* decode_ptr_host, int bbox_element, cv::Mat& img, int count) {
+void process_decode_ptr_host(std::vector& 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 &res, const float* decode_pt
}
}
}
-void batch_process(std::vector> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector& img_batch) {
+void batch_process(std::vector>& res_batch, const float* decode_ptr_host, int batch_size,
+ int bbox_element, const std::vector& img_batch) {
res_batch.resize(batch_size);
int count = static_cast(*decode_ptr_host);
count = count > kMaxNumOutputBbox ? kMaxNumOutputBbox : count;
@@ -225,4 +234,3 @@ void batch_process(std::vector> &res_batch, const float*
process_decode_ptr_host(res_batch[i], &decode_ptr_host[i * count], bbox_element, img, count);
}
}
-