haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Yolov8 seg (#1381)

Browse Source 
* v8_seg

* fix code style

* fix code style

* [优化代码风格]

* readme中增加yolov8_seg分割的使用

* 更新readme

* 更新readme

* [fix python infer bug]

* Update README.md

* Update README.md

---------

Co-authored-by: Wang Xinyu <shaywxy@gmail.com>
This commit is contained in:
Rex 2023-10-17 20:59:21 +08:00 committed by GitHub
parent 7b95f981f6
commit 9a53b7f226
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 1285 additions and 663 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
yolov8/CMakeLists.txt
View File

@ -41,10 +41,12 @@ include_directories(${OpenCV_INCLUDE_DIRS})
file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu)
add_executable(yolov8 ${PROJECT_SOURCE_DIR}/main.cpp ${SRCS})
add_executable(yolov8_det ${PROJECT_SOURCE_DIR}/yolov8_det.cpp ${SRCS})
target_link_libraries(yolov8 nvinfer)
target_link_libraries(yolov8 cudart)
target_link_libraries(yolov8 myplugins)
target_link_libraries(yolov8 ${OpenCV_LIBS})
target_link_libraries(yolov8_det nvinfer)
target_link_libraries(yolov8_det cudart)
target_link_libraries(yolov8_det myplugins)
target_link_libraries(yolov8_det ${OpenCV_LIBS})
add_executable(yolov8_seg ${PROJECT_SOURCE_DIR}/yolov8_seg.cpp ${SRCS})
target_link_libraries(yolov8_seg nvinfer cudart myplugins ${OpenCV_LIBS})

25
yolov8/README.md
View File

@ -9,7 +9,7 @@ The tensorrt code is derived from [xiaocao-tian/yolov8_tensorrt](https://github.
<a href="https://github.com/xiaocao-tian"><img src="https://avatars.githubusercontent.com/u/65889782?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/xinsuinizhuan"><img src="https://avatars.githubusercontent.com/u/40679769?v=4?s=48" width="40px;" alt=""/></a>
<a href="https://github.com/Rex-LK"><img src="https://avatars.githubusercontent.com/u/74702576?s=48&v=4" width="40px;" alt=""/></a>
## Requirements
@ -40,7 +40,7 @@ python gen_wts.py
```
2. build tensorrtx/yolov8 and run
### Detection
```
cd {tensorrtx}/yolov8/
// update kNumClass in config.h if your model is trained on custom dataset
@ -49,13 +49,24 @@ cd build
cp {ultralytics}/ultralytics/yolov8.wts {tensorrtx}/yolov8/build
cmake ..
make
sudo ./yolov8 -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file
sudo ./yolov8 -d [.engine] [image folder] [c/g] // deserialize and run inference, the images in [image folder] will be processed.
sudo ./yolov8_det -s [.wts] [.engine] [n/s/m/l/x] // 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 -s yolov8n.wts yolov8.engine n
sudo ./yolov8 -d yolov8n.engine ../images c //cpu postprocess
sudo ./yolov8 -d yolov8n.engine ../images g //gpu postprocess
sudo ./yolov8_det -s yolov8n.wts yolov8.engine n
sudo ./yolov8_det -d yolov8n.engine ../images c //cpu postprocess
sudo ./yolov8_det -d yolov8n.engine ../images g //gpu postprocess
```
### Instance Segmentation
```
# Build and serialize TensorRT engine
./yolov8_seg -s yolov8s-seg.wts yolov8s-seg.engine s
# Download the labels file
wget -O coco.txt https://raw.githubusercontent.com/amikelive/coco-labels/master/coco-labels-2014_2017.txt
# Run inference with labels file
./yolov8_seg -d yolov8s-seg.engine ../images c coco.txt //cpu postprocess
```
3. check the images generated, as follows. _zidane.jpg and _bus.jpg

2
yolov8/include/block.h
View File

@ -18,4 +18,4 @@ 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::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets);
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation = false);

17
yolov8/include/model.h
View File

@ -3,17 +3,8 @@
#include <string>
#include <assert.h>
nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
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* buildEngineYolov8s(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
nvinfer1::IHostMemory* buildEngineYolov8m(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
nvinfer1::IHostMemory* buildEngineYolov8l(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
nvinfer1::IHostMemory* buildEngineYolov8x(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path);
nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw, int& max_channels);

1
yolov8/include/postprocess.h
View File

@ -20,3 +20,4 @@ void cuda_decode(float* predict, int num_bboxes, float confidence_threshold,floa
void cuda_nms(float* parray, float nms_threshold, int max_objects, cudaStream_t stream);
void draw_mask_bbox(cv::Mat& img, std::vector<Detection>& dets, std::vector<cv::Mat>& masks, std::unordered_map<int, std::string>& labels_map);

1
yolov8/include/types.h
View File

@ -6,6 +6,7 @@ struct alignas(float) Detection {
float bbox[4];
float conf; // bbox_conf * cls_conf
float class_id;
float mask[32];
};
struct AffineMatrix {

39
yolov8/include/utils.h
View File

@ -1,6 +1,7 @@
#pragma once
#include <opencv2/opencv.hpp>
#include <dirent.h>
#include <fstream>
static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
int w, h, x, y;
@ -45,3 +46,41 @@ 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(' ');
if (std::string::npos == first) {
return str;
}
size_t last = str.find_last_not_of(' ');
return str.substr(first, (last - first + 1));
}
// Src: https://stackoverflow.com/questions/16605967
static inline std::string to_string_with_precision(const float a_value, const int n = 2) {
std::ostringstream out;
out.precision(n);
out << std::fixed << a_value;
return out.str();
}
static inline int read_labels(const std::string labels_filename, std::unordered_map<int, std::string>& labels_map) {
std::ifstream file(labels_filename);
// Read each line of the file
std::string line;
int index = 0;
while (std::getline(file, line)) {
// Strip the line of any leading or trailing whitespace
line = trim_leading_whitespace(line);
// Add the stripped line to the labels_map, using the loop index as the key
labels_map[index] = line;
index++;
}
// Close the file
file.close();
return 0;
}

23
yolov8/plugin/yololayer.cu
View File

@ -22,11 +22,12 @@ namespace Tn {
namespace nvinfer1 {
YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut) {
YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation) {
mClassCount = classCount;
mYoloV8NetWidth = netWidth;
mYoloV8netHeight = netHeight;
mMaxOutObject = maxOut;
is_segmentation_ = is_segmentation;
}
YoloLayerPlugin::~YoloLayerPlugin() {}
@ -39,6 +40,7 @@ YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) {
read(d, mYoloV8NetWidth);
read(d, mYoloV8netHeight);
read(d, mMaxOutObject);
read(d, is_segmentation_);
assert(d == a + length);
}
@ -52,12 +54,13 @@ void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT {
write(d, mYoloV8NetWidth);
write(d, mYoloV8netHeight);
write(d, mMaxOutObject);
write(d, is_segmentation_);
assert(d == a + getSerializationSize());
}
size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT {
return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mYoloV8netHeight) + sizeof(mYoloV8NetWidth) + sizeof(mMaxOutObject);
return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mYoloV8netHeight) + sizeof(mYoloV8NetWidth) + sizeof(mMaxOutObject) + sizeof(is_segmentation_);
}
int YoloLayerPlugin::initialize() TRT_NOEXCEPT {
@ -113,7 +116,7 @@ void YoloLayerPlugin::destroy() TRT_NOEXCEPT {
nvinfer1::IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT {
YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV8NetWidth, mYoloV8netHeight, mMaxOutObject);
YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV8NetWidth, mYoloV8netHeight, mMaxOutObject, is_segmentation_);
p->setPluginNamespace(mPluginNamespace);
return p;
}
@ -128,12 +131,13 @@ int YoloLayerPlugin::enqueue(int batchSize, const void* TRT_CONST_ENQUEUE* input
__device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); };
__global__ void CalDetection(const float* input, float* output, int numElements, int maxoutobject,
const int grid_h, int grid_w, const int stride, int classes, int outputElem) {
const int grid_h, int grid_w, const int stride, int classes, int outputElem, bool is_segmentation) {
int idx = threadIdx.x + blockDim.x * blockIdx.x;
if (idx >= numElements) return;
int total_grid = grid_h * grid_w;
int info_len = 4 + classes;
if (is_segmentation) info_len += 32;
int batchIdx = idx / total_grid;
int elemIdx = idx % total_grid;
const float* curInput = input + batchIdx * total_grid * info_len;
@ -141,7 +145,7 @@ __global__ void CalDetection(const float* input, float* output, int numElements,
int class_id = 0;
float max_cls_prob = 0.0;
for (int i = 4; i < info_len; i++) {
for (int i = 4; i < 4 + classes; i++) {
float p = Logist(curInput[elemIdx + i * total_grid]);
if (p > max_cls_prob) {
max_cls_prob = p;
@ -165,6 +169,10 @@ __global__ void CalDetection(const float* input, float* output, int numElements,
det->bbox[1] = (row + 0.5f - curInput[elemIdx + 1 * total_grid]) * stride;
det->bbox[2] = (col + 0.5f + curInput[elemIdx + 2 * total_grid]) * stride;
det->bbox[3] = (row + 0.5f + curInput[elemIdx + 3 * total_grid]) * stride;
for (int k = 0; is_segmentation && k < 32; k++) {
det->mask[k] = curInput[elemIdx + (k + 4 + classes) * total_grid];
}
}
void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,int mYoloV8NetWidth, int batchSize) {
@ -184,7 +192,7 @@ void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cuda
if (numElem < mThreadCount) mThreadCount = numElem;
CalDetection << <(numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> >
(inputs[i], output, numElem, mMaxOutObject, grid_h, grid_w, stride, mClassCount, outputElem);
(inputs[i], output, numElem, mMaxOutObject, grid_h, grid_w, stride, mClassCount, outputElem, is_segmentation_);
}
}
@ -217,7 +225,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];
YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count);
bool is_segmentation = p_netinfo[4];
YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, is_segmentation);
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}

3
yolov8/plugin/yololayer.h
View File

@ -7,7 +7,7 @@
namespace nvinfer1 {
class API YoloLayerPlugin : public IPluginV2IOExt {
public:
YoloLayerPlugin(int classCount, int netWdith, int netHeight, int maxOut);
YoloLayerPlugin(int classCount, int netWdith, int netHeight, int maxOut, bool is_segmentation);
YoloLayerPlugin(const void* data, size_t length);
~YoloLayerPlugin();
@ -66,6 +66,7 @@ public:
int mYoloV8NetWidth;
int mYoloV8netHeight;
int mMaxOutObject;
bool is_segmentation_;
};
class API YoloPluginCreator : public IPluginCreator {

6
yolov8/src/block.cpp
View File

@ -169,13 +169,13 @@ nvinfer1::ITensor& input, int ch, int grid, int k, int s, int p, std::string lna
}
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets) {
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation) {
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
nvinfer1::PluginField plugin_fields[1];
int netinfo[4] = {kNumClass, kInputW, kInputH, kMaxNumOutputBbox};
int netinfo[5] = {kNumClass, kInputW, kInputH, kMaxNumOutputBbox, is_segmentation};
plugin_fields[0].data = netinfo;
plugin_fields[0].length = 4;
plugin_fields[0].length = 5;
plugin_fields[0].name = "netinfo";
plugin_fields[0].type = nvinfer1::PluginFieldType::kFLOAT32;

816
yolov8/src/model.cpp
View File

@ -1,11 +1,67 @@
#include <math.h>
#include <iostream>
#include "model.h"
#include "block.h"
#include "calibrator.h"
#include <iostream>
#include "config.h"
nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path) {
static int get_width(int x, float gw, int max_channels, int divisor = 8) {
auto channel = int(ceil((x * gw) / divisor)) * divisor;
return channel >= max_channels ? max_channels : channel;
}
static int get_depth(int x, float gd) {
if (x == 1) return 1;
int r = round(x * gd);
if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) --r;
return std::max<int>(r, 1);
}
static nvinfer1::IElementWiseLayer* Proto(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap,
nvinfer1::ITensor& input, std::string lname, float gw, int max_channels) {
int mid_channel = get_width(256, gw, max_channels);
auto cv1 = convBnSiLU(network, weightMap, input, mid_channel, 3, 1, 1, "model.22.proto.cv1");
float* convTranpsose_bais = (float*)weightMap["model.22.proto.upsample.bias"].values;
int convTranpsose_bais_len = weightMap["model.22.proto.upsample.bias"].count;
nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, convTranpsose_bais, convTranpsose_bais_len};
auto convTranpsose = network->addDeconvolutionNd(*cv1->getOutput(0), mid_channel, nvinfer1::DimsHW{2,2}, weightMap["model.22.proto.upsample.weight"], bias);
assert(convTranpsose);
convTranpsose->setStrideNd(nvinfer1::DimsHW{2, 2});
auto cv2 = convBnSiLU(network,weightMap,*convTranpsose->getOutput(0), mid_channel, 3, 1, 1, "model.22.proto.cv2");
auto cv3 = convBnSiLU(network,weightMap,*cv2->getOutput(0), 32, 1, 1, 0,"model.22.proto.cv3");
assert(cv3);
return cv3;
}
static nvinfer1::IShuffleLayer* ProtoCoef(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap,
nvinfer1::ITensor& input, std::string lname, int grid_shape, float gw) {
int mid_channle = 0;
if(gw == 0.25 || gw== 0.5) {
mid_channle = 32;
} else if(gw == 0.75) {
mid_channle = 48;
} else if(gw == 1.00) {
mid_channle = 64;
} else if(gw == 1.25) {
mid_channle = 80;
}
auto cv0 = convBnSiLU(network, weightMap, input, mid_channle, 3, 1, 1, lname + ".0");
auto cv1 = convBnSiLU(network, weightMap, *cv0->getOutput(0), mid_channle, 3, 1, 1, lname + ".1");
float* cv2_bais_value = (float*)weightMap[lname + ".2" + ".bias"].values;
int cv2_bais_len = weightMap[lname + ".2" + ".bias"].count;
nvinfer1::Weights cv2_bais{nvinfer1::DataType::kFLOAT, cv2_bais_value, cv2_bais_len};
auto cv2 = network->addConvolutionNd(*cv1->getOutput(0), 32, nvinfer1::DimsHW{1, 1}, weightMap[lname + ".2" + ".weight"], cv2_bais);
cv2->setStrideNd(nvinfer1::DimsHW{1, 1});
nvinfer1::IShuffleLayer* cv2_shuffle = network->addShuffle(*cv2->getOutput(0));
cv2_shuffle->setReshapeDimensions(nvinfer1::Dims2{ 32, grid_shape});
return cv2_shuffle;
}
nvinfer1::IHostMemory* buildEngineYolov8Det(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
const std::string& wts_path, float& gd, float& gw, int& max_channels) {
std::map<std::string, nvinfer1::Weights> weightMap = loadWeights(wts_path);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(0U);
@ -18,16 +74,20 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder,
/*******************************************************************************************************
***************************************** YOLOV8 BACKBONE ********************************************
*******************************************************************************************************/
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, 16, 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), 32, 3, 2, 1, "model.1");
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), 32, 32, 1, true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), 64, 3, 2, 1, "model.3");
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), 64, 64, 2, true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), 128, 3, 2, 1, "model.5");
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), 128, 128, 2, true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), 256, 3, 2, 1, "model.7");
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), 256, 256, 1, true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), 256, 256, 5, "model.9");
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, get_width(64, gw, max_channels), 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), get_width(128, gw, max_channels), 3, 2, 1, "model.1");
// 11233
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), get_width(128, gw, max_channels), get_width(128, gw, max_channels), get_depth(3, gd), true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), get_width(256, gw, max_channels), 3, 2, 1, "model.3");
// 22466
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), get_width(256, gw, max_channels), get_width(256, gw, max_channels), get_depth(6, gd), true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), get_width(512, gw, max_channels), 3, 2, 1, "model.5");
// 22466
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), get_width(512, gw, max_channels), get_width(512, gw, max_channels), get_depth(6, gd), true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), get_width(1024, gw, max_channels), 3, 2, 1, "model.7");
// 11233
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), get_width(1024, gw, max_channels), get_width(1024, gw, max_channels), get_depth(3, gd), true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), get_width(1024, gw, max_channels), get_width(1024, gw, max_channels), 5, "model.9");
/*******************************************************************************************************
********************************************* YOLOV8 HEAD ********************************************
@ -40,8 +100,7 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder,
nvinfer1::ITensor* inputTensor11[] = {upsample10->getOutput(0), conv6->getOutput(0)};
nvinfer1::IConcatenationLayer* cat11 = network->addConcatenation(inputTensor11, 2);
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), 128, 128, 1, false, 0.5, "model.12");
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), get_width(512, gw, max_channels), get_width(512, gw, max_channels), get_depth(3, gd), false, 0.5, "model.12");
nvinfer1::IResizeLayer* upsample13 = network->addResize(*conv12->getOutput(0));
assert(upsample13);
@ -50,100 +109,95 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder,
nvinfer1::ITensor* inputTensor14[] = {upsample13->getOutput(0), conv4->getOutput(0)};
nvinfer1::IConcatenationLayer* cat14 = network->addConcatenation(inputTensor14, 2);
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), 64, 64, 1, false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 2, 1, "model.16");
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), get_width(256, gw, max_channels), get_width(256, gw, max_channels), get_depth(3, gd), false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), get_width(256, gw, max_channels), 3, 2, 1, "model.16");
nvinfer1::ITensor* inputTensor17[] = {conv16->getOutput(0), conv12->getOutput(0)};
nvinfer1::IConcatenationLayer* cat17 = network->addConcatenation(inputTensor17, 2);
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), 128, 128, 1, false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 128, 3, 2, 1, "model.19");
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), get_width(512, gw, max_channels), get_width(512, gw, max_channels), get_depth(3, gd), false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), get_width(512, gw, max_channels), 3, 2, 1, "model.19");
nvinfer1::ITensor* inputTensor20[] = {conv19->getOutput(0), conv9->getOutput(0)};
nvinfer1::IConcatenationLayer* cat20 = network->addConcatenation(inputTensor20, 2);
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), 256, 256, 1, false, 0.5, "model.21");
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), get_width(1024, gw, max_channels), get_width(1024, gw, max_channels), get_depth(3, gd), false, 0.5, "model.21");
/*******************************************************************************************************
********************************************* YOLOV8 OUTPUT ******************************************
*******************************************************************************************************/
// output0
int base_in_channel = (gw == 1.25) ? 80 : 64;
int base_out_channel = (gw == 0.25) ? 320 : 256;
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
// output0
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{1, 1});
conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{0, 0});
nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 64, 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
conv22_cv3_0_2->setStride(nvinfer1::DimsHW{1, 1});
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{0, 0});
nvinfer1::ITensor* inputTensor22_0[] = {conv22_cv2_0_2->getOutput(0), conv22_cv3_0_2->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_0 = network->addConcatenation(inputTensor22_0, 2);
// output1
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IConvolutionLayer* conv22_cv2_1_2 = network->addConvolutionNd(*conv22_cv2_1_1->getOutput(0), 64, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv2.1.2.weight"], weightMap["model.22.cv2.1.2.bias"]);
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{1,1});
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{0,0});
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 64, 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 64, 3, 1, 1, "model.22.cv3.1.1");
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{1, 1});
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{0, 0});
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.1.1");
nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]);
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{1,1});
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{0,0});
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{1, 1});
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{0, 0});
nvinfer1::ITensor* inputTensor22_1[] = {conv22_cv2_1_2->getOutput(0), conv22_cv3_1_2->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_1 = network->addConcatenation(inputTensor22_1, 2);
// output2
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 64, 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::ITensor* inputTensor22_2[] = {conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_2 = network->addConcatenation(inputTensor22_2, 2);
/*******************************************************************************************************
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 8) * (kInputW / 8) });
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 8) * (kInputW / 8)});
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{1,1});
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{ kNumClass, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{1,1});
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 8) * (kInputW / 8)}, nvinfer1::Dims2{1, 1});
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{kNumClass, (kInputH / 8) * (kInputW / 8)}, nvinfer1::Dims2{1, 1});
nvinfer1::IShuffleLayer* dfl22_0 = DFL(network, weightMap, *split22_0_0->getOutput(0), 4, (kInputH / 8) * (kInputW / 8), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_0[] = {dfl22_0->getOutput(0), split22_0_1->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_dfl_0 = network->addConcatenation(inputTensor22_dfl_0, 2);
nvinfer1::IShuffleLayer* shuffle22_1 = network->addShuffle(*cat22_1->getOutput(0));
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 16) * (kInputW / 16) });
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{1,1});
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{ kNumClass, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{1,1});
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 16) * (kInputW / 16)});
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 16) * (kInputW / 16)}, nvinfer1::Dims2{1, 1});
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{kNumClass, (kInputH / 16) * (kInputW / 16)}, nvinfer1::Dims2{1, 1});
nvinfer1::IShuffleLayer* dfl22_1 = DFL(network, weightMap, *split22_1_0->getOutput(0), 4, (kInputH / 16) * (kInputW / 16), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_1[] = {dfl22_1->getOutput(0), split22_1_1->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_dfl_1 = network->addConcatenation(inputTensor22_dfl_1, 2);
nvinfer1::IShuffleLayer* shuffle22_2 = network->addShuffle(*cat22_2->getOutput(0));
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 32) * (kInputW / 32) });
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{1,1});
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{ kNumClass, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{1,1});
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 32) * (kInputW / 32)});
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 32) * (kInputW / 32)}, nvinfer1::Dims2{1, 1});
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{kNumClass, (kInputH / 32) * (kInputW / 32)}, nvinfer1::Dims2{1, 1});
nvinfer1::IShuffleLayer* dfl22_2 = DFL(network, weightMap, *split22_2_0->getOutput(0), 4, (kInputH / 32) * (kInputW / 32), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_2[] = {dfl22_2->getOutput(0), split22_2_1->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_dfl_2 = network->addConcatenation(inputTensor22_dfl_2, 2);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer*>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2});
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer *>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2});
yolo->getOutput(0)->setName(kOutputTensorName);
network->markOutput(*yolo->getOutput(0));
builder->setMaxBatchSize(kBatchSize);
config->setMaxWorkspaceSize(16* (1<<20));
config->setMaxWorkspaceSize(16 * (1 << 20));
#if defined(USE_FP16)
config->setFlag(nvinfer1::BuilderFlag::kFP16);
@ -161,150 +215,159 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder,
delete network;
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
for (auto &mem : weightMap){
free((void *)(mem.second.values));
}
return serialized_model;
}
nvinfer1::IHostMemory* buildEngineYolov8s(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path) {
nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt,
const std::string& wts_path, float& gd, float& gw, int& max_channels) {
std::map<std::string, nvinfer1::Weights> weightMap = loadWeights(wts_path);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(0U);
/*******************************************************************************************************
****************************************** YOLOV8 INPUT **********************************************
*******************************************************************************************************/
nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt, nvinfer1::Dims3{ 3, kInputH, kInputW });
nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt, nvinfer1::Dims3{3, kInputH, kInputW});
assert(data);
/*******************************************************************************************************
***************************************** YOLOV8 BACKBONE ********************************************
*******************************************************************************************************/
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, 32, 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), 64, 3, 2, 1, "model.1");
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), 64, 64, 1, true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), 128, 3, 2, 1, "model.3");
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), 128, 128, 2, true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), 256, 3, 2, 1, "model.5");
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), 256, 256, 2, true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), 512, 3, 2, 1, "model.7");
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), 512, 512, 1, true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), 512, 512, 5, "model.9");
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, get_width(64, gw, max_channels), 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), get_width(128, gw, max_channels), 3, 2, 1, "model.1");
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), get_width(128, gw, max_channels), get_width(128, gw, max_channels), get_depth(3, gd), true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), get_width(256, gw, max_channels), 3, 2, 1, "model.3");
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), get_width(256, gw, max_channels), get_width(256, gw, max_channels), get_depth(6, gd), true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), get_width(512, gw, max_channels), 3, 2, 1, "model.5");
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), get_width(512, gw, max_channels), get_width(512, gw, max_channels), get_depth(6, gd), true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), get_width(1024, gw, max_channels), 3, 2, 1, "model.7");
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), get_width(1024, gw, max_channels), get_width(1024, gw, max_channels), get_depth(3, gd), true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), get_width(1024, gw, max_channels), get_width(1024, gw, max_channels), 5, "model.9");
/*******************************************************************************************************
********************************************* YOLOV8 HEAD ********************************************
*******************************************************************************************************/
float scale[] = { 1.0, 2.0, 2.0 };
float scale[] = {1.0, 2.0, 2.0};
nvinfer1::IResizeLayer* upsample10 = network->addResize(*conv9->getOutput(0));
assert(upsample10);
upsample10->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample10->setScales(scale, 3);
nvinfer1::ITensor* inputTensor11[] = { upsample10->getOutput(0), conv6->getOutput(0) };
nvinfer1::ITensor* inputTensor11[] = {upsample10->getOutput(0), conv6->getOutput(0)};
nvinfer1::IConcatenationLayer* cat11 = network->addConcatenation(inputTensor11, 2);
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), 256, 256, 1, false, 0.5, "model.12");
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), get_width(512, gw, max_channels), get_width(512, gw, max_channels), get_depth(3, gd), false, 0.5, "model.12");
nvinfer1::IResizeLayer* upsample13 = network->addResize(*conv12->getOutput(0));
assert(upsample13);
upsample13->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample13->setScales(scale, 3);
nvinfer1::ITensor* inputTensor14[] = { upsample13->getOutput(0), conv4->getOutput(0) };
nvinfer1::ITensor* inputTensor14[] = {upsample13->getOutput(0), conv4->getOutput(0)};
nvinfer1::IConcatenationLayer* cat14 = network->addConcatenation(inputTensor14, 2);
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), 128, 128, 1, false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 128, 3, 2, 1, "model.16");
nvinfer1::ITensor* inputTensor17[] = { conv16->getOutput(0), conv12->getOutput(0) };
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), get_width(256, gw, max_channels), get_width(256, gw, max_channels), get_depth(3, gd), false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), get_width(256, gw, max_channels), 3, 2, 1, "model.16");
nvinfer1::ITensor* inputTensor17[] = {conv16->getOutput(0), conv12->getOutput(0)};
nvinfer1::IConcatenationLayer* cat17 = network->addConcatenation(inputTensor17, 2);
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), 256, 256, 1, false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 256, 3, 2, 1, "model.19");
nvinfer1::ITensor* inputTensor20[] = { conv19->getOutput(0), conv9->getOutput(0) };
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), get_width(512, gw, max_channels), get_width(512, gw, max_channels), get_depth(3, gd), false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), get_width(512, gw, max_channels), 3, 2, 1, "model.19");
nvinfer1::ITensor* inputTensor20[] = {conv19->getOutput(0), conv9->getOutput(0)};
nvinfer1::IConcatenationLayer* cat20 = network->addConcatenation(inputTensor20, 2);
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), 512, 512, 1, false, 0.5, "model.21");
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), get_width(1024, gw, max_channels), get_width(1024, gw, max_channels), get_depth(3, gd), false, 0.5, "model.21");
/*******************************************************************************************************
********************************************* YOLOV8 OUTPUT ******************************************
*******************************************************************************************************/
int base_in_channel = (gw == 1.25) ? 80 : 64;
int base_out_channel = (gw == 0.25) ? 320 : 256;
// output0
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{ 1, 1 });
conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 128, 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 128, 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
conv22_cv3_0_2->setStride(nvinfer1::DimsHW{ 1, 1 });
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::ITensor* inputTensor22_0[] = { conv22_cv2_0_2->getOutput(0), conv22_cv3_0_2->getOutput(0) };
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{1, 1});
conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{0, 0});
nvinfer1::IElementWiseLayer *conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer *conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer *conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
conv22_cv3_0_2->setStride(nvinfer1::DimsHW{1, 1});
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{0, 0});
nvinfer1::ITensor* inputTensor22_0[] = {conv22_cv2_0_2->getOutput(0), conv22_cv3_0_2->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_0 = network->addConcatenation(inputTensor22_0, 2);
// output1
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IConvolutionLayer* conv22_cv2_1_2 = network->addConvolutionNd(*conv22_cv2_1_1->getOutput(0), 64, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv2.1.2.weight"], weightMap["model.22.cv2.1.2.bias"]);
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 128, 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 128, 3, 1, 1, "model.22.cv3.1.1");
nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]);
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::ITensor* inputTensor22_1[] = { conv22_cv2_1_2->getOutput(0), conv22_cv3_1_2->getOutput(0) };
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IConvolutionLayer* conv22_cv2_1_2 = network->addConvolutionNd(*conv22_cv2_1_1->getOutput(0), 64, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv2.1.2.weight"], weightMap["model.22.cv2.1.2.bias"]);
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{1, 1});
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{0, 0});
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.1.1");
nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]);
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{1, 1});
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{0, 0});
nvinfer1::ITensor* inputTensor22_1[] = {conv22_cv2_1_2->getOutput(0), conv22_cv3_1_2->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_1 = network->addConcatenation(inputTensor22_1, 2);
// output2
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 128, 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 128, 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::ITensor* inputTensor22_2[] = { conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0) };
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), base_in_channel, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), get_width(base_out_channel, gw, max_channels), 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::ITensor* inputTensor22_2[] = {conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0)};
nvinfer1::IConcatenationLayer* cat22_2 = network->addConcatenation(inputTensor22_2, 2);
/*******************************************************************************************************
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 8) * (kInputW / 8) });
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 8) * (kInputW / 8)});
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 8) * (kInputW / 8)}, nvinfer1::Dims2{1, 1});
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{kNumClass, (kInputH / 8) * (kInputW / 8)}, nvinfer1::Dims2{1, 1});
nvinfer1::IShuffleLayer* dfl22_0 = DFL(network, weightMap, *split22_0_0->getOutput(0), 4, (kInputH / 8) * (kInputW / 8), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_0[] = { dfl22_0->getOutput(0), split22_0_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_0 = network->addConcatenation(inputTensor22_dfl_0, 2);
nvinfer1::IShuffleLayer* shuffle22_1 = network->addShuffle(*cat22_1->getOutput(0));
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 16) * (kInputW / 16) });
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 16) * (kInputW / 16)});
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 16) * (kInputW / 16)}, nvinfer1::Dims2{1, 1});
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{kNumClass, (kInputH / 16) * (kInputW / 16)}, nvinfer1::Dims2{1, 1});
nvinfer1::IShuffleLayer* dfl22_1 = DFL(network, weightMap, *split22_1_0->getOutput(0), 4, (kInputH / 16) * (kInputW / 16), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_1[] = { dfl22_1->getOutput(0), split22_1_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_1 = network->addConcatenation(inputTensor22_dfl_1, 2);
nvinfer1::IShuffleLayer* shuffle22_2 = network->addShuffle(*cat22_2->getOutput(0));
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 32) * (kInputW / 32) });
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass, (kInputH / 32) * (kInputW / 32)});
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{0, 0}, nvinfer1::Dims2{64, (kInputH / 32) * (kInputW / 32)}, nvinfer1::Dims2{1, 1});
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{64, 0}, nvinfer1::Dims2{kNumClass, (kInputH / 32) * (kInputW / 32)}, nvinfer1::Dims2{1, 1});
nvinfer1::IShuffleLayer* dfl22_2 = DFL(network, weightMap, *split22_2_0->getOutput(0), 4, (kInputH / 32) * (kInputW / 32), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_2[] = { dfl22_2->getOutput(0), split22_2_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_2 = network->addConcatenation(inputTensor22_dfl_2, 2);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer*>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2});
// det0
auto proto_coef_0 = ProtoCoef(network, weightMap, *conv15->getOutput(0), "model.22.cv4.0", 6400, gw);
nvinfer1::ITensor* inputTensor22_dfl_0[] = { dfl22_0->getOutput(0), split22_0_1->getOutput(0),proto_coef_0->getOutput(0)};
nvinfer1::IConcatenationLayer *cat22_dfl_0 = network->addConcatenation(inputTensor22_dfl_0, 3);
// det1
auto proto_coef_1 = ProtoCoef(network, weightMap, *conv18->getOutput(0), "model.22.cv4.1", 1600, gw);
nvinfer1::ITensor* inputTensor22_dfl_1[] = { dfl22_1->getOutput(0), split22_1_1->getOutput(0),proto_coef_1->getOutput(0)};
nvinfer1::IConcatenationLayer *cat22_dfl_1 = network->addConcatenation(inputTensor22_dfl_1, 3);
// det2
auto proto_coef_2 = ProtoCoef(network, weightMap, *conv21->getOutput(0), "model.22.cv4.2", 400, gw);
nvinfer1::ITensor* inputTensor22_dfl_2[] = { dfl22_2->getOutput(0), split22_2_1->getOutput(0) ,proto_coef_2->getOutput(0)};
nvinfer1::IConcatenationLayer *cat22_dfl_2 = network->addConcatenation(inputTensor22_dfl_2, 3);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer *>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2}, true);
yolo->getOutput(0)->setName(kOutputTensorName);
network->markOutput(*yolo->getOutput(0));
auto proto = Proto(network, weightMap, *conv15->getOutput(0), "model.22.proto", gw, max_channels);
proto->getOutput(0)->setName("proto");
network->markOutput(*proto->getOutput(0));
builder->setMaxBatchSize(kBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20));
@ -329,470 +392,3 @@ nvinfer1::IHostMemory* buildEngineYolov8s(nvinfer1::IBuilder* builder,
}
return serialized_model;
}
nvinfer1::IHostMemory* buildEngineYolov8m(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path) {
std::map<std::string, nvinfer1::Weights> weightMap = loadWeights(wts_path);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(0U);
/*******************************************************************************************************
****************************************** YOLOV8 INPUT **********************************************
*******************************************************************************************************/
nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt, nvinfer1::Dims3{ 3, kInputH, kInputW });
assert(data);
/*******************************************************************************************************
***************************************** YOLOV8 BACKBONE ********************************************
*******************************************************************************************************/
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, 48, 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), 96, 3, 2, 1, "model.1");
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), 96, 96, 2, true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), 192, 3, 2, 1, "model.3");
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), 192, 192, 4, true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), 384, 3, 2, 1, "model.5");
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), 384, 384, 4, true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), 576, 3, 2, 1, "model.7");
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), 576, 576, 2, true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), 576, 576, 5, "model.9");
/*******************************************************************************************************
********************************************* YOLOV8 HEAD ********************************************
*******************************************************************************************************/
float scale[] = { 1.0, 2.0, 2.0 };
nvinfer1::IResizeLayer* upsample10 = network->addResize(*conv9->getOutput(0));
upsample10->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample10->setScales(scale, 3);
nvinfer1::ITensor* inputTensor11[] = { upsample10->getOutput(0), conv6->getOutput(0) };
nvinfer1::IConcatenationLayer* cat11 = network->addConcatenation(inputTensor11, 2);
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), 384, 384, 2, false, 0.5, "model.12");
nvinfer1::IResizeLayer* upsample13 = network->addResize(*conv12->getOutput(0));
upsample13->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample13->setScales(scale, 3);
nvinfer1::ITensor* inputTensor14[] = { upsample13->getOutput(0), conv4->getOutput(0) };
nvinfer1::IConcatenationLayer* cat14 = network->addConcatenation(inputTensor14, 2);
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), 192, 192, 2, false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 192, 3, 2, 1, "model.16");
nvinfer1::ITensor* inputTensor17[] = { conv16->getOutput(0), conv12->getOutput(0) };
nvinfer1::IConcatenationLayer* cat17 = network->addConcatenation(inputTensor17, 2);
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), 384, 384, 2, false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 384, 3, 2, 1, "model.19");
nvinfer1::ITensor* inputTensor20[] = { conv19->getOutput(0), conv9->getOutput(0) };
nvinfer1::IConcatenationLayer* cat20 = network->addConcatenation(inputTensor20, 2);
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), 576, 576, 2, false, 0.5, "model.21");
/*******************************************************************************************************
********************************************* YOLOV8 OUTPUT ******************************************
*******************************************************************************************************/
// output0
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{ 1, 1 });
conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 192, 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 192, 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
conv22_cv3_0_2->setStride(nvinfer1::DimsHW{ 1, 1 });
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::ITensor* inputTensor22_0[] = { conv22_cv2_0_2->getOutput(0), conv22_cv3_0_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_0 = network->addConcatenation(inputTensor22_0, 2);
// output1
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IConvolutionLayer* conv22_cv2_1_2 = network->addConvolutionNd(*conv22_cv2_1_1->getOutput(0), 64, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv2.1.2.weight"], weightMap["model.22.cv2.1.2.bias"]);
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 192, 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 192, 3, 1, 1, "model.22.cv3.1.1");
nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]);
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::ITensor* inputTensor22_1[] = { conv22_cv2_1_2->getOutput(0), conv22_cv3_1_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_1 = network->addConcatenation(inputTensor22_1, 2);
// output2
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 192, 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 192, 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::ITensor* inputTensor22_2[] = { conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_2 = network->addConcatenation(inputTensor22_2, 2);
/*******************************************************************************************************
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 8) * (kInputW / 8) });
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_0 = DFL(network, weightMap, *split22_0_0->getOutput(0), 4, (kInputH / 8) * (kInputW / 8), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_0[] = { dfl22_0->getOutput(0), split22_0_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_0 = network->addConcatenation(inputTensor22_dfl_0, 2);
nvinfer1::IShuffleLayer* shuffle22_1 = network->addShuffle(*cat22_1->getOutput(0));
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 16) * (kInputW / 16) });
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_1 = DFL(network, weightMap, *split22_1_0->getOutput(0), 4, (kInputH / 16) * (kInputW / 16), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_1[] = { dfl22_1->getOutput(0), split22_1_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_1 = network->addConcatenation(inputTensor22_dfl_1, 2);
nvinfer1::IShuffleLayer* shuffle22_2 = network->addShuffle(*cat22_2->getOutput(0));
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 32) * (kInputW / 32) });
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_2 = DFL(network, weightMap, *split22_2_0->getOutput(0), 4, (kInputH / 32) * (kInputW / 32), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_2[] = { dfl22_2->getOutput(0), split22_2_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_2 = network->addConcatenation(inputTensor22_dfl_2, 2);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer*>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2});
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);
nvinfer1::IInt8EntropyCalibrator2* calibrator = new Calibrator(1, kInputW, kInputH, "../calibrator/", "int8calib.table", kInputTensorName);
config->setInt8Calibrator(calibrator);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
nvinfer1::IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
delete network;
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return serialized_model;
}
nvinfer1::IHostMemory* buildEngineYolov8l(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path) {
std::map<std::string, nvinfer1::Weights> weightMap = loadWeights(wts_path);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(0U);
/*******************************************************************************************************
****************************************** YOLOV8 INPUT **********************************************
*******************************************************************************************************/
nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt, nvinfer1::Dims3{ 3, kInputH, kInputW });
assert(data);
/*******************************************************************************************************
***************************************** YOLOV8 BACKBONE ********************************************
*******************************************************************************************************/
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), 128, 3, 2, 1, "model.1");
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), 128, 128, 3, true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), 256, 3, 2, 1, "model.3");
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), 256, 256, 6, true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), 512, 3, 2, 1, "model.5");
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), 512, 512, 6, true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), 512, 3, 2, 1, "model.7");
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), 512, 512, 3, true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), 512, 512, 5, "model.9");
/*******************************************************************************************************
****************************************** YOLOV8 HEAD ***********************************************
*******************************************************************************************************/
float scale[] = { 1.0, 2.0, 2.0 };
nvinfer1::IResizeLayer* upsample10 = network->addResize(*conv9->getOutput(0));
upsample10->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample10->setScales(scale, 3);
nvinfer1::ITensor* inputTensor11[] = { upsample10->getOutput(0), conv6->getOutput(0) };
nvinfer1::IConcatenationLayer* cat11 = network->addConcatenation(inputTensor11, 2);
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), 512, 512, 3, false, 0.5, "model.12");
nvinfer1::IResizeLayer* upsample13 = network->addResize(*conv12->getOutput(0));
upsample13->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample13->setScales(scale, 3);
nvinfer1::ITensor* inputTensor14[] = { upsample13->getOutput(0), conv4->getOutput(0) };
nvinfer1::IConcatenationLayer* cat14 = network->addConcatenation(inputTensor14, 2);
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), 256, 256, 3, false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 256, 3, 2, 1, "model.16");
nvinfer1::ITensor* inputTensor17[] = { conv16->getOutput(0), conv12->getOutput(0) };
nvinfer1::IConcatenationLayer* cat17 = network->addConcatenation(inputTensor17, 2);
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), 512, 512, 3, false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 512, 3, 2, 1, "model.19");
nvinfer1::ITensor* inputTensor20[] = { conv19->getOutput(0), conv9->getOutput(0) };
nvinfer1::IConcatenationLayer* cat20 = network->addConcatenation(inputTensor20, 2);
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), 512, 512, 3, false, 0.5, "model.21");
/*******************************************************************************************************
********************************************* YOLOV8 OUTPUT ******************************************
*******************************************************************************************************/
// output0
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{ 1, 1 });
conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 256, 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 256, 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
conv22_cv3_0_2->setStride(nvinfer1::DimsHW{ 1, 1 });
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::ITensor* inputTensor22_0[] = { conv22_cv2_0_2->getOutput(0), conv22_cv3_0_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_0 = network->addConcatenation(inputTensor22_0, 2);
// output1
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IConvolutionLayer* conv22_cv2_1_2 = network->addConvolutionNd(*conv22_cv2_1_1->getOutput(0), 64, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv2.1.2.weight"], weightMap["model.22.cv2.1.2.bias"]);
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 256, 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 256, 3, 1, 1, "model.22.cv3.1.1");
nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]);
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::ITensor* inputTensor22_1[] = { conv22_cv2_1_2->getOutput(0), conv22_cv3_1_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_1 = network->addConcatenation(inputTensor22_1, 2);
// output2
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 256, 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 256, 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::ITensor* inputTensor22_2[] = { conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_2 = network->addConcatenation(inputTensor22_2, 2);
/*******************************************************************************************************
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 8) * (kInputW / 8) });
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_0 = DFL(network, weightMap, *split22_0_0->getOutput(0), 4, (kInputH / 8) * (kInputW / 8), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_0[] = { dfl22_0->getOutput(0), split22_0_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_0 = network->addConcatenation(inputTensor22_dfl_0, 2);
nvinfer1::IShuffleLayer* shuffle22_1 = network->addShuffle(*cat22_1->getOutput(0));
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 16) * (kInputW / 16) });
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_1 = DFL(network, weightMap, *split22_1_0->getOutput(0), 4, (kInputH / 16) * (kInputW / 16), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_1[] = { dfl22_1->getOutput(0), split22_1_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_1 = network->addConcatenation(inputTensor22_dfl_1, 2);
nvinfer1::IShuffleLayer* shuffle22_2 = network->addShuffle(*cat22_2->getOutput(0));
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 32) * (kInputW / 32) });
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_2 = DFL(network, weightMap, *split22_2_0->getOutput(0), 4, (kInputH / 32) * (kInputW / 32), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_2[] = { dfl22_2->getOutput(0), split22_2_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_2 = network->addConcatenation(inputTensor22_dfl_2, 2);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer*>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2});
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);
nvinfer1::IInt8EntropyCalibrator2* calibrator = new Calibrator(1, kInputW, kInputH, "../calibrator/", "int8calib.table", kInputTensorName);
config->setInt8Calibrator(calibrator);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
nvinfer1::IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
delete network;
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return serialized_model;
}
nvinfer1::IHostMemory* buildEngineYolov8x(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path) {
std::map<std::string, nvinfer1::Weights> weightMap = loadWeights(wts_path);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(0U);
/*******************************************************************************************************
****************************************** YOLOV8 INPUT **********************************************
*******************************************************************************************************/
nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt, nvinfer1::Dims3{ 3, kInputH, kInputW });
assert(data);
/*******************************************************************************************************
***************************************** YOLOV8 BACKBONE ********************************************
*******************************************************************************************************/
nvinfer1::IElementWiseLayer* conv0 = convBnSiLU(network, weightMap, *data, 80, 3, 2, 1, "model.0");
nvinfer1::IElementWiseLayer* conv1 = convBnSiLU(network, weightMap, *conv0->getOutput(0), 160, 3, 2, 1, "model.1");
nvinfer1::IElementWiseLayer* conv2 = C2F(network, weightMap, *conv1->getOutput(0), 160, 160, 3, true, 0.5, "model.2");
nvinfer1::IElementWiseLayer* conv3 = convBnSiLU(network, weightMap, *conv2->getOutput(0), 320, 3, 2, 1, "model.3");
nvinfer1::IElementWiseLayer* conv4 = C2F(network, weightMap, *conv3->getOutput(0), 320, 320, 6, true, 0.5, "model.4");
nvinfer1::IElementWiseLayer* conv5 = convBnSiLU(network, weightMap, *conv4->getOutput(0), 640, 3, 2, 1, "model.5");
nvinfer1::IElementWiseLayer* conv6 = C2F(network, weightMap, *conv5->getOutput(0), 640, 640, 6, true, 0.5, "model.6");
nvinfer1::IElementWiseLayer* conv7 = convBnSiLU(network, weightMap, *conv6->getOutput(0), 640, 3, 2, 1, "model.7");
nvinfer1::IElementWiseLayer* conv8 = C2F(network, weightMap, *conv7->getOutput(0), 640, 640, 3, true, 0.5, "model.8");
nvinfer1::IElementWiseLayer* conv9 = SPPF(network, weightMap, *conv8->getOutput(0), 640, 640, 5, "model.9");
/*******************************************************************************************************
****************************************** YOLOV8 HEAD ***********************************************
*******************************************************************************************************/
float scale[] = { 1.0, 2.0, 2.0 };
nvinfer1::IResizeLayer* upsample10 = network->addResize(*conv9->getOutput(0));
upsample10->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample10->setScales(scale, 3);
nvinfer1::ITensor* inputTensor11[] = { upsample10->getOutput(0), conv6->getOutput(0) };
nvinfer1::IConcatenationLayer* cat11 = network->addConcatenation(inputTensor11, 2);
nvinfer1::IElementWiseLayer* conv12 = C2F(network, weightMap, *cat11->getOutput(0), 640, 640, 3, false, 0.5, "model.12");
nvinfer1::IResizeLayer* upsample13 = network->addResize(*conv12->getOutput(0));
upsample13->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample13->setScales(scale, 3);
nvinfer1::ITensor* inputTensor14[] = { upsample13->getOutput(0), conv4->getOutput(0) };
nvinfer1::IConcatenationLayer* cat14 = network->addConcatenation(inputTensor14, 2);
nvinfer1::IElementWiseLayer* conv15 = C2F(network, weightMap, *cat14->getOutput(0), 320, 320, 3, false, 0.5, "model.15");
nvinfer1::IElementWiseLayer* conv16 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 320, 3, 2, 1, "model.16");
nvinfer1::ITensor* inputTensor17[] = { conv16->getOutput(0), conv12->getOutput(0) };
nvinfer1::IConcatenationLayer* cat17 = network->addConcatenation(inputTensor17, 2);
nvinfer1::IElementWiseLayer* conv18 = C2F(network, weightMap, *cat17->getOutput(0), 640, 640, 3, false, 0.5, "model.18");
nvinfer1::IElementWiseLayer* conv19 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 640, 3, 2, 1, "model.19");
nvinfer1::ITensor* inputTensor20[] = { conv19->getOutput(0), conv9->getOutput(0) };
nvinfer1::IConcatenationLayer* cat20 = network->addConcatenation(inputTensor20, 2);
nvinfer1::IElementWiseLayer* conv21 = C2F(network, weightMap, *cat20->getOutput(0), 640, 640, 3, false, 0.5, "model.21");
/*******************************************************************************************************
********************************************* YOLOV8 OUTPUT ******************************************
*******************************************************************************************************/
// output0
nvinfer1::IElementWiseLayer* conv22_cv2_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 80, 3, 1, 1, "model.22.cv2.0.0");
nvinfer1::IElementWiseLayer* conv22_cv2_0_1 = convBnSiLU(network, weightMap, *conv22_cv2_0_0->getOutput(0), 80, 3, 1, 1, "model.22.cv2.0.1");
nvinfer1::IConvolutionLayer* conv22_cv2_0_2 = network->addConvolutionNd(*conv22_cv2_0_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.0.2.weight"], weightMap["model.22.cv2.0.2.bias"]);
conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{ 1, 1 });
conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 320, 3, 1, 1, "model.22.cv3.0.0");
nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 320, 3, 1, 1, "model.22.cv3.0.1");
nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]);
conv22_cv3_0_2->setStride(nvinfer1::DimsHW{ 1, 1 });
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{ 0, 0 });
nvinfer1::ITensor* inputTensor22_0[] = { conv22_cv2_0_2->getOutput(0), conv22_cv3_0_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_0 = network->addConcatenation(inputTensor22_0, 2);
// output1
nvinfer1::IElementWiseLayer* conv22_cv2_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 80, 3, 1, 1, "model.22.cv2.1.0");
nvinfer1::IElementWiseLayer* conv22_cv2_1_1 = convBnSiLU(network, weightMap, *conv22_cv2_1_0->getOutput(0), 80, 3, 1, 1, "model.22.cv2.1.1");
nvinfer1::IConvolutionLayer* conv22_cv2_1_2 = network->addConvolutionNd(*conv22_cv2_1_1->getOutput(0), 64, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv2.1.2.weight"], weightMap["model.22.cv2.1.2.bias"]);
conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 320, 3, 1, 1, "model.22.cv3.1.0");
nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 320, 3, 1, 1, "model.22.cv3.1.1");
nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1, 1 }, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]);
conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{ 1,1 });
conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{ 0,0 });
nvinfer1::ITensor* inputTensor22_1[] = { conv22_cv2_1_2->getOutput(0), conv22_cv3_1_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_1 = network->addConcatenation(inputTensor22_1, 2);
// output2
nvinfer1::IElementWiseLayer* conv22_cv2_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 80, 3, 1, 1, "model.22.cv2.2.0");
nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), 80, 3, 1, 1, "model.22.cv2.2.1");
nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]);
nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 320, 3, 1, 1, "model.22.cv3.2.0");
nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 320, 3, 1, 1, "model.22.cv3.2.1");
nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{ 1,1 }, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]);
nvinfer1::ITensor* inputTensor22_2[] = { conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_2 = network->addConcatenation(inputTensor22_2, 2);
/*******************************************************************************************************
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));
shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 8) * (kInputW / 8) });
nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_0_1 = network->addSlice(*shuffle22_0->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 8) * (kInputW / 8) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_0 = DFL(network, weightMap, *split22_0_0->getOutput(0), 4, (kInputH / 8) * (kInputW / 8), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_0[] = { dfl22_0->getOutput(0), split22_0_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_0 = network->addConcatenation(inputTensor22_dfl_0, 2);
nvinfer1::IShuffleLayer* shuffle22_1 = network->addShuffle(*cat22_1->getOutput(0));
shuffle22_1->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 16) * (kInputW / 16) });
nvinfer1::ISliceLayer* split22_1_0 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_1_1 = network->addSlice(*shuffle22_1->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 16) * (kInputW / 16) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_1 = DFL(network, weightMap, *split22_1_0->getOutput(0), 4, (kInputH / 16) * (kInputW / 16), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_1[] = { dfl22_1->getOutput(0), split22_1_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_1 = network->addConcatenation(inputTensor22_dfl_1, 2);
nvinfer1::IShuffleLayer* shuffle22_2 = network->addShuffle(*cat22_2->getOutput(0));
shuffle22_2->setReshapeDimensions(nvinfer1::Dims2{ 64 + kNumClass, (kInputH / 32) * (kInputW / 32) });
nvinfer1::ISliceLayer* split22_2_0 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 0, 0 }, nvinfer1::Dims2{ 64, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::ISliceLayer* split22_2_1 = network->addSlice(*shuffle22_2->getOutput(0), nvinfer1::Dims2{ 64, 0 }, nvinfer1::Dims2{ kNumClass, (kInputH / 32) * (kInputW / 32) }, nvinfer1::Dims2{ 1,1 });
nvinfer1::IShuffleLayer* dfl22_2 = DFL(network, weightMap, *split22_2_0->getOutput(0), 4, (kInputH / 32) * (kInputW / 32), 1, 1, 0, "model.22.dfl.conv.weight");
nvinfer1::ITensor* inputTensor22_dfl_2[] = { dfl22_2->getOutput(0), split22_2_1->getOutput(0) };
nvinfer1::IConcatenationLayer* cat22_dfl_2 = network->addConcatenation(inputTensor22_dfl_2, 2);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, std::vector<nvinfer1::IConcatenationLayer*>{cat22_dfl_0, cat22_dfl_1, cat22_dfl_2});
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);
nvinfer1::IInt8EntropyCalibrator2* calibrator = new Calibrator(1, kInputW, kInputH, "../calibrator/", "int8calib.table", kInputTensorName);
config->setInt8Calibrator(calibrator);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
nvinfer1::IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
delete network;
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return serialized_model;
}

66
yolov8/src/postprocess.cpp
View File

@ -1,5 +1,5 @@
#include "postprocess.h"
#include "utils.h"
cv::Rect get_rect(cv::Mat &img, float bbox[4]) {
float l, r, t, b;
@ -121,3 +121,67 @@ void draw_bbox(std::vector<cv::Mat> &img_batch, std::vector<std::vector<Detectio
}
}
}
cv::Mat scale_mask(cv::Mat mask, cv::Mat img) {
int x, y, w, h;
float r_w = kInputW / (img.cols * 1.0);
float r_h = kInputH / (img.rows * 1.0);
if (r_h > r_w) {
w = kInputW;
h = r_w * img.rows;
x = 0;
y = (kInputH - h) / 2;
} else {
w = r_h * img.cols;
h = kInputH;
x = (kInputW - w) / 2;
y = 0;
}
cv::Rect r(x, y, w, h);
cv::Mat res;
cv::resize(mask(r), res, img.size());
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};
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()];
auto bgr = cv::Scalar(color & 0xFF, color >> 8 & 0xFF, color >> 16 & 0xFF);
cv::Rect r = get_rect(img, dets[i].bbox);
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;
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;
}
}
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);
// Set the top left corner of the rectangle
cv::Point topLeft(r.x, r.y - textSize.height);
// Set the bottom right corner of the rectangle
cv::Point bottomRight(r.x + textSize.width, r.y + textSize.height);
// Set the thickness of the rectangle lines
int lineThickness = 2;
// 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);
}
}

46
yolov8/main.cpp → yolov8/yolov8_det.cpp
View File

@ -13,22 +13,12 @@ Logger gLogger;
using namespace nvinfer1;
const int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1;
void serialize_engine(std::string &wts_name, std::string &engine_name, std::string &sub_type) {
void serialize_engine(std::string &wts_name, std::string &engine_name, std::string &sub_type, float &gd, float &gw, int &max_channels) {
IBuilder *builder = createInferBuilder(gLogger);
IBuilderConfig *config = builder->createBuilderConfig();
IHostMemory *serialized_engine = nullptr;
if (sub_type == "n") {
serialized_engine = buildEngineYolov8n(builder, config, DataType::kFLOAT, wts_name);
} else if (sub_type == "s") {
serialized_engine = buildEngineYolov8s(builder, config, DataType::kFLOAT, wts_name);
} else if (sub_type == "m") {
serialized_engine = buildEngineYolov8m(builder, config, DataType::kFLOAT, wts_name);
} else if (sub_type == "l") {
serialized_engine = buildEngineYolov8l(builder, config, DataType::kFLOAT, wts_name);
} else if (sub_type == "x") {
serialized_engine = buildEngineYolov8x(builder, config, DataType::kFLOAT, wts_name);
}
serialized_engine = buildEngineYolov8Det(builder, config, DataType::kFLOAT, wts_name, gd, gw, max_channels);
assert(serialized_engine);
std::ofstream p(engine_name, std::ios::binary);
@ -114,12 +104,36 @@ void infer(IExecutionContext &context, cudaStream_t &stream, void **buffers, flo
}
bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, std::string &img_dir, std::string &sub_type, std::string &cuda_post_process) {
bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, std::string &img_dir, std::string &sub_type,
std::string &cuda_post_process, float &gd, float &gw, int &max_channels) {
if (argc < 4) return false;
if (std::string(argv[1]) == "-s" && argc == 5) {
wts = std::string(argv[2]);
engine = std::string(argv[3]);
sub_type = std::string(argv[4]);
if (sub_type == "n") {
gd = 0.33;
gw = 0.25;
max_channels = 1024;
} else if (sub_type == "s"){
gd = 0.33;
gw = 0.50;
max_channels = 1024;
} else if (sub_type == "m") {
gd = 0.67;
gw = 0.75;
max_channels = 576;
} else if (sub_type == "l") {
gd = 1.0;
gw = 1.0;
max_channels = 512;
} else if (sub_type == "x") {
gd = 1.0;
gw = 1.25;
max_channels = 640;
} else {
return false;
}
} else if (std::string(argv[1]) == "-d" && argc == 5) {
engine = std::string(argv[2]);
img_dir = std::string(argv[3]);
@ -138,8 +152,10 @@ int main(int argc, char **argv) {
std::string sub_type = "";
std::string cuda_post_process="";
int model_bboxes;
float gd = 0.0f, gw = 0.0f;
int max_channels = 0;
if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type, cuda_post_process)) {
if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type, cuda_post_process, gd, gw, max_channels)) {
std::cerr << "Arguments not right!" << std::endl;
std::cerr << "./yolov8 -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file" << std::endl;
std::cerr << "./yolov8 -d [.engine] ../samples [c/g]// deserialize plan file and run inference" << std::endl;
@ -148,7 +164,7 @@ int main(int argc, char **argv) {
// Create a model using the API directly and serialize it to a file
if (!wts_name.empty()) {
serialize_engine(wts_name, engine_name, sub_type);
serialize_engine(wts_name, engine_name, sub_type, gd, gw, max_channels);
return 0;
}

0
yolov8/yolov8_trt.py → yolov8/yolov8_det_trt.py
View File

321
yolov8/yolov8_seg.cpp Normal file
View File

@ -0,0 +1,321 @@
#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "model.h"
#include "utils.h"
#include "preprocess.h"
#include "postprocess.h"
#include "cuda_utils.h"
#include "logging.h"
Logger gLogger;
using namespace nvinfer1;
const int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1;
const static int kOutputSegSize = 32 * (kInputH / 4) * (kInputW / 4);
static cv::Rect get_downscale_rect(float bbox[4], float scale) {
float left = bbox[0];
float top = bbox[1];
float right = bbox[0] + bbox[2];
float bottom = bbox[1] + bbox[3];
left = left < 0 ? 0 : left;
top = top < 0 ? 0: top;
right = right > 640 ? 640 : right;
bottom = bottom > 640 ? 640: bottom;
left /= scale;
top /= scale;
right /= scale;
bottom /= scale;
return cv::Rect(int(left), int(top), int(right - left), int(bottom - top));
}
std::vector<cv::Mat> process_mask(const float* proto, int proto_size, std::vector<Detection>& dets) {
std::vector<cv::Mat> masks;
for (size_t i = 0; i < dets.size(); i++) {
cv::Mat mask_mat = cv::Mat::zeros(kInputH / 4, kInputW / 4, CV_32FC1);
auto r = get_downscale_rect(dets[i].bbox, 4);
for (int x = r.x; x < r.x + r.width; x++) {
for (int y = r.y; y < r.y + r.height; y++) {
float e = 0.0f;
for (int j = 0; j < 32; j++) {
e += dets[i].mask[j] * proto[j * proto_size / 32 + y * mask_mat.cols + x];
}
e = 1.0f / (1.0f + expf(-e));
mask_mat.at<float>(y, x) = e;
}
}
cv::resize(mask_mat, mask_mat, cv::Size(kInputW, kInputH));
masks.push_back(mask_mat);
}
return masks;
}
void serialize_engine(std::string &wts_name, std::string &engine_name, std::string &sub_type, float &gd, float &gw, int &max_channels)
{
IBuilder *builder = createInferBuilder(gLogger);
IBuilderConfig *config = builder->createBuilderConfig();
IHostMemory *serialized_engine = nullptr;
serialized_engine = buildEngineYolov8Seg(builder, config, DataType::kFLOAT, wts_name, gd, gw, max_channels);
assert(serialized_engine);
std::ofstream p(engine_name, std::ios::binary);
if (!p)
{
std::cout << "could not open plan output file" << std::endl;
assert(false);
}
p.write(reinterpret_cast<const char *>(serialized_engine->data()), serialized_engine->size());
delete builder;
delete config;
delete serialized_engine;
}
void deserialize_engine(std::string &engine_name, IRuntime **runtime, ICudaEngine **engine, IExecutionContext **context)
{
std::ifstream file(engine_name, std::ios::binary);
if (!file.good())
{
std::cerr << "read " << engine_name << " error!" << std::endl;
assert(false);
}
size_t size = 0;
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
char *serialized_engine = new char[size];
assert(serialized_engine);
file.read(serialized_engine, size);
file.close();
*runtime = createInferRuntime(gLogger);
assert(*runtime);
*engine = (*runtime)->deserializeCudaEngine(serialized_engine, size);
assert(*engine);
*context = (*engine)->createExecutionContext();
assert(*context);
delete[] serialized_engine;
}
void prepare_buffer(ICudaEngine *engine, float **input_buffer_device, float **output_buffer_device, float **output_seg_buffer_device,
float **output_buffer_host,float **output_seg_buffer_host ,float **decode_ptr_host, float **decode_ptr_device, std::string cuda_post_process) {
assert(engine->getNbBindings() == 3);
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine->getBindingIndex(kInputTensorName);
const int outputIndex = engine->getBindingIndex(kOutputTensorName);
const int outputIndex_seg = engine->getBindingIndex("proto");
assert(inputIndex == 0);
assert(outputIndex == 1);
assert(outputIndex_seg == 2);
// Create GPU buffers on device
CUDA_CHECK(cudaMalloc((void **) input_buffer_device, kBatchSize * 3 * kInputH * kInputW * sizeof(float)));
CUDA_CHECK(cudaMalloc((void **) output_buffer_device, kBatchSize * kOutputSize * sizeof(float)));
CUDA_CHECK(cudaMalloc((void **) output_seg_buffer_device, kBatchSize * kOutputSegSize * sizeof(float)));
if (cuda_post_process == "c") {
*output_buffer_host = new float[kBatchSize * kOutputSize];
*output_seg_buffer_host = new float[kBatchSize * kOutputSegSize];
} else if (cuda_post_process == "g") {
if (kBatchSize > 1) {
std::cerr << "Do not yet support GPU post processing for multiple batches" << std::endl;
exit(0);
}
// Allocate memory for decode_ptr_host and copy to device
*decode_ptr_host = new float[1 + kMaxNumOutputBbox * bbox_element];
CUDA_CHECK(cudaMalloc((void **)decode_ptr_device, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element)));
}
}
void infer(IExecutionContext &context, cudaStream_t &stream, void **buffers, float *output, float *output_seg,int batchsize, float* decode_ptr_host, float* decode_ptr_device, int model_bboxes, std::string cuda_post_process) {
// infer on the batch asynchronously, and DMA output back to host
auto start = std::chrono::system_clock::now();
context.enqueue(batchsize, buffers, stream, nullptr);
if (cuda_post_process == "c") {
std::cout << "kOutputSize:" << kOutputSize <<std::endl;
CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchsize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost,stream));
std::cout << "kOutputSegSize:" << kOutputSegSize <<std::endl;
CUDA_CHECK(cudaMemcpyAsync(output_seg, buffers[2], batchsize * kOutputSegSize * sizeof(float), cudaMemcpyDeviceToHost, stream));
auto end = std::chrono::system_clock::now();
std::cout << "inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
} else if (cuda_post_process == "g") {
CUDA_CHECK(cudaMemsetAsync(decode_ptr_device, 0, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element), stream));
cuda_decode((float *)buffers[1], model_bboxes, kConfThresh, decode_ptr_device, kMaxNumOutputBbox, stream);
cuda_nms(decode_ptr_device, kNmsThresh, kMaxNumOutputBbox, stream);//cuda nms
CUDA_CHECK(cudaMemcpyAsync(decode_ptr_host, decode_ptr_device, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element), cudaMemcpyDeviceToHost, stream));
auto end = std::chrono::system_clock::now();
std::cout << "inference and gpu postprocess time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
CUDA_CHECK(cudaStreamSynchronize(stream));
}
bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, std::string &img_dir, std::string &sub_type,
std::string &cuda_post_process, std::string labels_filename, float &gd, float &gw, int &max_channels)
{
if (argc < 4)
return false;
if (std::string(argv[1]) == "-s" && argc == 5) {
wts = std::string(argv[2]);
engine = std::string(argv[3]);
sub_type = std::string(argv[4]);
if (sub_type == "n") {
gd = 0.33;
gw = 0.25;
max_channels = 1024;
} else if (sub_type == "s") {
gd = 0.33;
gw = 0.50;
max_channels = 1024;
} else if (sub_type == "m") {
gd = 0.67;
gw = 0.75;
max_channels = 576;
} else if (sub_type == "l") {
gd = 1.0;
gw = 1.0;
max_channels = 512;
} else if (sub_type == "x") {
gd = 1.0;
gw = 1.25;
max_channels = 640;
} else{
return false;
}
} else if (std::string(argv[1]) == "-d" && argc == 6) {
engine = std::string(argv[2]);
img_dir = std::string(argv[3]);
cuda_post_process = std::string(argv[4]);
labels_filename = std::string(argv[5]);
} else {
return false;
}
return true;
}
int main(int argc, char **argv) {
cudaSetDevice(kGpuId);
std::string wts_name = "";
std::string engine_name = "";
std::string img_dir;
std::string sub_type = "";
std::string cuda_post_process = "";
std::string labels_filename = "../coco.txt";
int model_bboxes;
float gd = 0.0f, gw = 0.0f;
int max_channels = 0;
if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type, cuda_post_process, labels_filename, gd, gw, max_channels)) {
std::cerr << "Arguments not right!" << std::endl;
std::cerr << "./yolov8 -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file" << std::endl;
std::cerr << "./yolov8 -d [.engine] ../samples [c/g] coco_file// deserialize plan file and run inference" << std::endl;
return -1;
}
// Create a model using the API directly and serialize it to a file
if (!wts_name.empty()) {
serialize_engine(wts_name, engine_name, sub_type, gd, gw, max_channels);
return 0;
}
// Deserialize the engine from file
IRuntime *runtime = nullptr;
ICudaEngine *engine = nullptr;
IExecutionContext *context = nullptr;
deserialize_engine(engine_name, &runtime, &engine, &context);
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
cuda_preprocess_init(kMaxInputImageSize);
auto out_dims = engine->getBindingDimensions(1);
model_bboxes = out_dims.d[0];
// Prepare cpu and gpu buffers
float *device_buffers[3];
float *output_buffer_host = nullptr;
float *output_seg_buffer_host = nullptr;
float *decode_ptr_host=nullptr;
float *decode_ptr_device=nullptr;
// Read images from directory
std::vector<std::string> file_names;
if (read_files_in_dir(img_dir.c_str(), file_names) < 0) {
std::cerr << "read_files_in_dir failed." << std::endl;
return -1;
}
std::unordered_map<int, std::string> labels_map;
read_labels(labels_filename, labels_map);
assert(kNumClass == labels_map.size());
prepare_buffer(engine, &device_buffers[0], &device_buffers[1], &device_buffers[2], &output_buffer_host, &output_seg_buffer_host,&decode_ptr_host, &decode_ptr_device, cuda_post_process);
// // batch predict
for (size_t i = 0; i < file_names.size(); i += kBatchSize) {
// Get a batch of images
std::vector<cv::Mat> img_batch;
std::vector<std::string> img_name_batch;
for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) {
cv::Mat img = cv::imread(img_dir + "/" + file_names[j]);
img_batch.push_back(img);
img_name_batch.push_back(file_names[j]);
}
// Preprocess
cuda_batch_preprocess(img_batch, device_buffers[0], kInputW, kInputH, stream);
// Run inference
infer(*context, stream, (void **)device_buffers, output_buffer_host, output_seg_buffer_host,kBatchSize, decode_ptr_host, decode_ptr_device, model_bboxes, cuda_post_process);
std::vector<std::vector<Detection>> res_batch;
if (cuda_post_process == "c") {
// NMS
batch_nms(res_batch, output_buffer_host, img_batch.size(), kOutputSize, kConfThresh, kNmsThresh);
for (size_t b = 0; b < img_batch.size(); b++) {
auto& res = res_batch[b];
cv::Mat img = img_batch[b];
auto masks = process_mask(&output_seg_buffer_host[b * kOutputSegSize], kOutputSegSize, res);
draw_mask_bbox(img, res, masks, labels_map);
cv::imwrite("_" + img_name_batch[b], img);
}
} else if (cuda_post_process == "g") {
// Process gpu decode and nms results
// batch_process(res_batch, decode_ptr_host, img_batch.size(), bbox_element, img_batch);
// todo seg in gpu
std::cerr << "seg_postprocess is not support in gpu right now" << std::endl;
}
}
// Release stream and buffers
cudaStreamDestroy(stream);
CUDA_CHECK(cudaFree(device_buffers[0]));
CUDA_CHECK(cudaFree(device_buffers[1]));
CUDA_CHECK(cudaFree(device_buffers[2]));
CUDA_CHECK(cudaFree(decode_ptr_device));
delete[] decode_ptr_host;
delete[] output_buffer_host;
delete[] output_seg_buffer_host;
cuda_preprocess_destroy();
// Destroy the engine
delete context;
delete engine;
delete runtime;
// Print histogram of the output distribution
// std::cout << "\nOutput:\n\n";
// for (unsigned int i = 0; i < kOutputSize; i++)
//{
// std::cout << prob[i] << ", ";
// if (i % 10 == 0) std::cout << std::endl;
//}
// std::cout << std::endl;
return 0;
}

570
yolov8/yolov8_seg_trt.py Normal file
View File

@ -0,0 +1,570 @@
"""
An example that uses TensorRT's Python api to make inferences.
"""
import ctypes
import os
import shutil
import random
import sys
import threading
import time
import cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
CONF_THRESH = 0.5
IOU_THRESHOLD = 0.4
def get_img_path_batches(batch_size, img_dir):
ret = []
batch = []
for root, dirs, files in os.walk(img_dir):
for name in files:
if len(batch) == batch_size:
ret.append(batch)
batch = []
batch.append(os.path.join(root, name))
if len(batch) > 0:
ret.append(batch)
return ret
def plot_one_box(x, img, color=None, label=None, line_thickness=None):
"""
description: Plots one bounding box on image img,
this function comes from YoLov8 project.
param:
x: a box likes [x1,y1,x2,y2]
img: a opencv image object
color: color to draw rectangle, such as (0,255,0)
label: str
line_thickness: int
return:
no return
"""
tl = (
line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
) # line/font thickness
color = color or [random.randint(0, 255) for _ in range(3)]
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
if label:
tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
cv2.putText(
img,
label,
(c1[0], c1[1] - 2),
0,
tl / 3,
[225, 255, 255],
thickness=tf,
lineType=cv2.LINE_AA,
)
class YoLov8TRT(object):
"""
description: A YOLOv8 class that warps TensorRT ops, preprocess and postprocess ops.
"""
def __init__(self, engine_file_path):
# Create a Context on this device,
self.ctx = cuda.Device(0).make_context()
stream = cuda.Stream()
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(TRT_LOGGER)
# Deserialize the engine from file
with open(engine_file_path, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
host_inputs = []
cuda_inputs = []
host_outputs = []
cuda_outputs = []
bindings = []
for binding in engine:
print('bingding:', binding, engine.get_binding_shape(binding))
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
cuda_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(cuda_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
self.input_w = engine.get_binding_shape(binding)[-1]
self.input_h = engine.get_binding_shape(binding)[-2]
host_inputs.append(host_mem)
cuda_inputs.append(cuda_mem)
else:
host_outputs.append(host_mem)
cuda_outputs.append(cuda_mem)
# Store
self.stream = stream
self.context = context
self.engine = engine
self.host_inputs = host_inputs
self.cuda_inputs = cuda_inputs
self.host_outputs = host_outputs
self.cuda_outputs = cuda_outputs
self.bindings = bindings
self.batch_size = engine.max_batch_size
#Data length
self.det_output_length = host_outputs[0].shape[0]
self.seg_output_length = host_outputs[1].shape[0]
self.seg_w = int(self.input_w / 4)
self.seg_h = int(self.input_h / 4)
self.seg_c = int(self.seg_output_length / (self.seg_w * self.seg_w))
self.det_row_output_length = self.seg_c + 6
# Draw mask
self.colors_obj = Colors()
def infer(self, raw_image_generator):
threading.Thread.__init__(self)
# Make self the active context, pushing it on top of the context stack.
self.ctx.push()
# Restore
stream = self.stream
context = self.context
engine = self.engine
host_inputs = self.host_inputs
cuda_inputs = self.cuda_inputs
host_outputs = self.host_outputs
cuda_outputs = self.cuda_outputs
bindings = self.bindings
# Do image preprocess
batch_image_raw = []
batch_origin_h = []
batch_origin_w = []
batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w])
for i, image_raw in enumerate(raw_image_generator):
input_image, image_raw, origin_h, origin_w = self.preprocess_image(image_raw)
batch_image_raw.append(image_raw)
batch_origin_h.append(origin_h)
batch_origin_w.append(origin_w)
np.copyto(batch_input_image[i], input_image)
batch_input_image = np.ascontiguousarray(batch_input_image)
# Copy input image to host buffer
np.copyto(host_inputs[0], batch_input_image.ravel())
start = time.time()
# Transfer input data to the GPU.
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
# Run inference.
context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
cuda.memcpy_dtoh_async(host_outputs[1], cuda_outputs[1], stream)
# Synchronize the stream
stream.synchronize()
end = time.time()
# Remove any context from the top of the context stack, deactivating it.
self.ctx.pop()
# Here we use the first row of output in that batch_size = 1
output = host_outputs[0]
output_proto_mask = host_outputs[1]
# Do postprocess
for i in range(self.batch_size):
result_boxes, result_scores, result_classid,result_proto_coef = self.post_process(
output[i * 38001: (i + 1) * 38001], batch_origin_h[i], batch_origin_w[i]
)
if result_proto_coef.shape[0] == 0:
continue
result_masks = self.process_mask(output_proto_mask, result_proto_coef, result_boxes, batch_origin_h[i], batch_origin_w[i])
self.draw_mask(result_masks, colors_=[self.colors_obj(x, True) for x in result_classid],im_src=batch_image_raw[i])
# Draw rectangles and labels on the original image
for j in range(len(result_boxes)):
box = result_boxes[j]
plot_one_box(
box,
batch_image_raw[i],
label="{}:{:.2f}".format(
categories[int(result_classid[j])], result_scores[j]
),
)
return batch_image_raw, end - start
def destroy(self):
# Remove any context from the top of the context stack, deactivating it.
self.ctx.pop()
def get_raw_image(self, image_path_batch):
"""
description: Read an image from image path
"""
for img_path in image_path_batch:
yield cv2.imread(img_path)
def get_raw_image_zeros(self, image_path_batch=None):
"""
description: Ready data for warmup
"""
for _ in range(self.batch_size):
yield np.zeros([self.input_h, self.input_w, 3], dtype=np.uint8)
def preprocess_image(self, raw_bgr_image):
"""
description: Convert BGR image to RGB,
resize and pad it to target size, normalize to [0,1],
transform to NCHW format.
param:
input_image_path: str, image path
return:
image: the processed image
image_raw: the original image
h: original height
w: original width
"""
image_raw = raw_bgr_image
h, w, c = image_raw.shape
image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)
# Calculate widht and height and paddings
r_w = self.input_w / w
r_h = self.input_h / h
if r_h > r_w:
tw = self.input_w
th = int(r_w * h)
tx1 = tx2 = 0
ty1 = int((self.input_h - th) / 2)
ty2 = self.input_h - th - ty1
else:
tw = int(r_h * w)
th = self.input_h
tx1 = int((self.input_w - tw) / 2)
tx2 = self.input_w - tw - tx1
ty1 = ty2 = 0
# Resize the image with long side while maintaining ratio
image = cv2.resize(image, (tw, th))
# Pad the short side with (128,128,128)
image = cv2.copyMakeBorder(
image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128)
)
image = image.astype(np.float32)
# Normalize to [0,1]
image /= 255.0
# HWC to CHW format:
image = np.transpose(image, [2, 0, 1])
# CHW to NCHW format
image = np.expand_dims(image, axis=0)
# Convert the image to row-major order, also known as "C order":
image = np.ascontiguousarray(image)
return image, image_raw, h, w
def xywh2xyxy(self, origin_h, origin_w, x):
"""
description: Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
param:
origin_h: height of original image
origin_w: width of original image
x: A boxes numpy, each row is a box [center_x, center_y, w, h]
return:
y: A boxes numpy, each row is a box [x1, y1, x2, y2]
"""
y = np.zeros_like(x)
r_w = self.input_w / origin_w
r_h = self.input_h / origin_h
if r_h > r_w:
y[:, 0] = x[:, 0]
y[:, 2] = x[:, 2]
y[:, 1] = x[:, 1] - (self.input_h - r_w * origin_h) / 2
y[:, 3] = x[:, 3] - (self.input_h - r_w * origin_h) / 2
y /= r_w
else:
y[:, 0] = x[:, 0] - (self.input_w - r_h * origin_w) / 2
y[:, 2] = x[:, 2] - (self.input_w - r_h * origin_w) / 2
y[:, 1] = x[:, 1]
y[:, 3] = x[:, 3]
y /= r_h
return y
def post_process(self, output, origin_h, origin_w):
"""
description: postprocess the prediction
param:
output: A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...]
origin_h: height of original image
origin_w: width of original image
return:
result_boxes: finally boxes, a boxes numpy, each row is a box [x1, y1, x2, y2]
result_scores: finally scores, a numpy, each element is the score correspoing to box
result_classid: finally classid, a numpy, each element is the classid correspoing to box
"""
# Get the num of boxes detected
num = int(output[0])
# Reshape to a two dimentional ndarray
pred = np.reshape(output[1:], (-1, 38))[:num, :]
# Do nms
boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD)
result_boxes = boxes[:, :4] if len(boxes) else np.array([])
result_scores = boxes[:, 4] if len(boxes) else np.array([])
result_classid = boxes[:, 5] if len(boxes) else np.array([])
result_proto_coef = boxes[:, 6:] if len(boxes) else np.array([])
return result_boxes, result_scores, result_classid,result_proto_coef
def bbox_iou(self, box1, box2, x1y1x2y2=True):
"""
description: compute the IoU of two bounding boxes
param:
box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
x1y1x2y2: select the coordinate format
return:
iou: computed iou
"""
if not x1y1x2y2:
# Transform from center and width to exact coordinates
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
else:
# Get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
# Get the coordinates of the intersection rectangle
inter_rect_x1 = np.maximum(b1_x1, b2_x1)
inter_rect_y1 = np.maximum(b1_y1, b2_y1)
inter_rect_x2 = np.minimum(b1_x2, b2_x2)
inter_rect_y2 = np.minimum(b1_y2, b2_y2)
# Intersection area
inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \
np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None)
# Union Area
b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)
iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)
return iou
def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4):
"""
description: Removes detections with lower object confidence score than 'conf_thres' and performs
Non-Maximum Suppression to further filter detections.
param:
prediction: detections, (x1, y1, x2, y2, conf, cls_id)
origin_h: original image height
origin_w: original image width
conf_thres: a confidence threshold to filter detections
nms_thres: a iou threshold to filter detections
return:
boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id)
"""
# Get the boxes that score > CONF_THRESH
boxes = prediction[prediction[:, 4] >= conf_thres]
# Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]
boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4])
# clip the coordinates
boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w - 1)
boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w - 1)
boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h - 1)
boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h - 1)
# Object confidence
confs = boxes[:, 4]
# Sort by the confs
boxes = boxes[np.argsort(-confs)]
# Perform non-maximum suppression
keep_boxes = []
while boxes.shape[0]:
large_overlap = self.bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres
label_match = boxes[0, 5] == boxes[:, 5]
# Indices of boxes with lower confidence scores, large IOUs and matching labels
invalid = large_overlap & label_match
keep_boxes += [boxes[0]]
boxes = boxes[~invalid]
boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([])
return boxes
def sigmoid(self, x):
return 1 / (1 + np.exp(-x))
def scale_mask(self, mask, ih, iw):
mask = cv2.resize(mask, (self.input_w, self.input_h))
r_w = self.input_w / (iw * 1.0)
r_h = self.input_h / (ih * 1.0)
if r_h > r_w:
w = self.input_w
h = int(r_w * ih)
x = 0
y = int((self.input_h - h) / 2)
else:
w = int(r_h * iw)
h = self.input_h
x = int((self.input_w - w) / 2)
y = 0
crop = mask[y:y+h, x:x+w]
crop = cv2.resize(crop, (iw, ih))
return crop
def process_mask(self, output_proto_mask, result_proto_coef, result_boxes, ih, iw):
"""
description: Mask pred by yolov8 instance segmentation ,
param:
output_proto_mask: prototype mask e.g. (32, 160, 160) for 640x640 input
result_proto_coef: prototype mask coefficients (n, 32), n represents n results
result_boxes :
ih: rows of original image
iw: cols of original image
return:
mask_result: (n, ih, iw)
"""
result_proto_masks = output_proto_mask.reshape(self.seg_c, self.seg_h, self.seg_w)
c, mh, mw = result_proto_masks.shape
masks = self.sigmoid((result_proto_coef @ result_proto_masks.astype(np.float32).reshape(c, -1))).reshape(-1, mh, mw)
mask_result = []
for mask, box in zip(masks, result_boxes):
mask_s = np.zeros((ih, iw))
crop_mask = self.scale_mask(mask, ih, iw)
x1 = int(box[0])
y1 = int(box[1])
x2 = int(box[2])
y2 = int(box[3])
crop = crop_mask[y1:y2, x1:x2]
crop = np.where(crop >= 0.5, 1, 0)
crop = crop.astype(np.uint8)
mask_s[y1:y2, x1:x2] = crop
mask_result.append(mask_s)
mask_result = np.array(mask_result)
return mask_result
def draw_mask(self, masks, colors_, im_src, alpha=0.5):
"""
description: Draw mask on image ,
param:
masks : result_mask
colors_: color to draw mask
im_src : original image
alpha : scale between original image and mask
return:
no return
"""
if len(masks) == 0:
return
masks = np.asarray(masks, dtype=np.uint8)
masks = np.ascontiguousarray(masks.transpose(1, 2, 0))
masks = np.asarray(masks, dtype=np.float32)
colors_ = np.asarray(colors_, dtype=np.float32)
s = masks.sum(2, keepdims=True).clip(0, 1)
masks = (masks @ colors_).clip(0, 255)
im_src[:] = masks * alpha + im_src * (1 - s * alpha)
class inferThread(threading.Thread):
def __init__(self, yolov8_wrapper, image_path_batch):
threading.Thread.__init__(self)
self.yolov8_wrapper = yolov8_wrapper
self.image_path_batch = image_path_batch
def run(self):
batch_image_raw, use_time = self.yolov8_wrapper.infer(self.yolov8_wrapper.get_raw_image(self.image_path_batch))
for i, img_path in enumerate(self.image_path_batch):
parent, filename = os.path.split(img_path)
save_name = os.path.join('output', filename)
# Save image
cv2.imwrite(save_name, batch_image_raw[i])
print('input->{}, time->{:.2f}ms, saving into output/'.format(self.image_path_batch, use_time * 1000))
class warmUpThread(threading.Thread):
def __init__(self, yolov8_wrapper):
threading.Thread.__init__(self)
self.yolov8_wrapper = yolov8_wrapper
def run(self):
batch_image_raw, use_time = self.yolov8_wrapper.infer(self.yolov8_wrapper.get_raw_image_zeros())
print('warm_up->{}, time->{:.2f}ms'.format(batch_image_raw[0].shape, use_time * 1000))
class Colors:
def __init__(self):
hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A',
'92CC17', '3DDB86', '1A9334', '00D4BB', '2C99A8', '00C2FF',
'344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF',
'FF95C8', 'FF37C7')
self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
self.n = len(self.palette)
def __call__(self, i, bgr=False):
c = self.palette[int(i) % self.n]
return (c[2], c[1], c[0]) if bgr else c
@staticmethod
def hex2rgb(h): # rgb order (PIL)
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
if __name__ == "__main__":
# load custom plugin and engine
PLUGIN_LIBRARY = "build/libmyplugins.so"
engine_file_path = "yolov8s-seg.engine"
if len(sys.argv) > 1:
engine_file_path = sys.argv[1]
if len(sys.argv) > 2:
PLUGIN_LIBRARY = sys.argv[2]
ctypes.CDLL(PLUGIN_LIBRARY)
# load coco labels
categories = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
"traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase",
"frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
"surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
"cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
"teddy bear",
"hair drier", "toothbrush"]
if os.path.exists('output/'):
shutil.rmtree('output/')
os.makedirs('output/')
# a YoLov8TRT instance
yolov8_wrapper = YoLov8TRT(engine_file_path)
try:
print('batch size is', yolov8_wrapper.batch_size)
image_dir = "images/"
image_path_batches = get_img_path_batches(yolov8_wrapper.batch_size, image_dir)
for i in range(10):
# create a new thread to do warm_up
thread1 = warmUpThread(yolov8_wrapper)
thread1.start()
thread1.join()
for batch in image_path_batches:
# create a new thread to do inference
thread1 = inferThread(yolov8_wrapper, batch)
thread1.start()
thread1.join()
finally:
# destroy the instance
yolov8_wrapper.destroy()