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

DBNet update (#318)

Browse Source 
dynamic input, optimize pre and post process
This commit is contained in:
BaofengZan 2020-12-04 11:38:12 +08:00 committed by GitHub
parent 53709d2ef6
commit 4089c64522
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 294 additions and 337 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
dbnet/README.md
View File

@ -3,9 +3,10 @@
The Pytorch implementation is [DBNet](https://github.com/BaofengZan/DBNet.pytorch).
<p align="center">
<img src="https://user-images.githubusercontent.com/20653176/89722330-00c36900-da1b-11ea-97f4-c61f9cd196fa.png">
<img src="https://user-images.githubusercontent.com/20653176/100968101-b044be00-356b-11eb-808c-9597cbe1f8de.jpg">
</p>
## How to Run
* 1. generate .wts
@ -34,7 +35,7 @@ https://github.com/BaofengZan/DBNet-TensorRT
## Todo
* 1. In common.hpp, the following two functions can be merged.
* 1. ~~In common.hpp, the following two functions can be merged.~~
```c++
ILayer* convBnLeaky(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, bool bias = true)
@ -46,4 +47,5 @@ ILayer* convBnLeaky2(INetworkDefinition *network, std::map<std::string, Weights>
* 2. The postprocess method here should be optimized, which is a little different from pytorch side.
* 3. The input image here is resized to 640x640 directly, while the pytorch side is using `letterbox` method.
* 3. ~~The input image here is resized to 640x640 directly, while the pytorch side is using `letterbox` method.~~

60
dbnet/common.hpp
View File

@ -39,9 +39,8 @@ std::map<std::string, Weights> loadWeights(const std::string file) {
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
while (count--) {
Weights wt{ DataType::kFLOAT, nullptr, 0 };
uint32_t size;
// Read name and type of blob
@ -51,12 +50,11 @@ std::map<std::string, Weights> loadWeights(const std::string file) {
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
for (uint32_t x = 0, y = size; x < y; ++x) {
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
@ -75,19 +73,19 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, W
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
Weights scale{ DataType::kFLOAT, scval, len };
float *shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
Weights shift{ DataType::kFLOAT, shval, len };
float *pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
Weights power{ DataType::kFLOAT, pval, len };
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
@ -97,52 +95,28 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, W
return scale_1;
}
ILayer* convBnLeaky(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, bool bias = true) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
ILayer* convBnLeaky(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, std::string bnname, bool bias = true) {
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
int p = ksize / 2;
IConvolutionLayer* conv1 = nullptr;
if (bias)
{
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], weightMap[lname + ".conv.bias"]);
if (bias) {
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".weight"], weightMap[lname + ".bias"]);
}
else
{
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts);
}
assert(conv1);
conv1->setStrideNd(DimsHW{s, s});
conv1->setPaddingNd(DimsHW{p, p});
conv1->setNbGroups(g);
//IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-4);
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3);
auto lr = network->addActivation(*bn1->getOutput(0), ActivationType::kLEAKY_RELU);
lr->setAlpha(0.1);
return lr;
}
ILayer* convBnLeaky2(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname, bool bias = true) {
Weights emptywts{DataType::kFLOAT, nullptr, 0};
int p = ksize / 2;
IConvolutionLayer* conv1 = nullptr;
if (bias)
{
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".0.weight"], weightMap[lname + ".0.bias"]);
}
else
{
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".0.weight"], emptywts);
else {
conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".weight"], emptywts);
}
assert(conv1);
conv1->setStrideNd(DimsHW{ s, s });
conv1->setPaddingNd(DimsHW{ p, p });
conv1->setNbGroups(g);
//IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-4);
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".1", 1e-3);
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname.substr(0, lname.find_last_of(".")) + bnname, 1e-3);
auto lr = network->addActivation(*bn1->getOutput(0), ActivationType::kLEAKY_RELU);
lr->setAlpha(0.1);
return lr;
}
IActivationLayer* basicBlock(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int stride, std::string lname) {
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
@ -187,7 +161,7 @@ int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_nam
struct dirent* p_file = nullptr;
while ((p_file = readdir(p_dir)) != nullptr) {
if (strcmp(p_file->d_name, ".") != 0 &&
strcmp(p_file->d_name, "..") != 0) {
strcmp(p_file->d_name, "..") != 0) {
//std::string cur_file_name(p_dir_name);
//cur_file_name += "/";
//cur_file_name += p_file->d_name;

563
dbnet/dbnet.cpp
View File

@ -3,49 +3,72 @@
#include "cuda_runtime_api.h"
#include "logging.h"
#include "common.hpp"
#include <math.h>
#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define BATCH_SIZE 1
#define EXPANDRATIO 1.4
static const int INPUT_H = 640;
static const int INPUT_W = 640;
static const int OUTPUT_SIZE = 640*640*2;
static const int SHORT_INPUT = 640;
static const int MAX_INPUT_SIZE = 1440; // 32x
static const int MIN_INPUT_SIZE = 608;
static const int OPT_INPUT_W = 1152;
static const int OPT_INPUT_H = 640;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "out";
static Logger gLogger;
cv::RotatedRect expandBox(const cv::RotatedRect& inBox, float ratio = 1.0)
{
cv::RotatedRect expandBox(const cv::RotatedRect& inBox, float ratio = 1.0) {
cv::Size size = inBox.size;
int neww = size.width * ratio;
int newh = size.height *ratio;
return cv::RotatedRect(inBox.center, cv::Size(neww, newh), inBox.angle);
}
float paddimg(cv::Mat& In_Out_img, int shortsize = 960) {
int w = In_Out_img.cols;
int h = In_Out_img.rows;
float scale = 1.f;
if (w < h) {
scale = (float)shortsize / w;
h = scale * h;
w = shortsize;
}
else {
scale = (float)shortsize / h;
w = scale * w;
h = shortsize;
}
if (h % 32 != 0) {
h = (h / 32 + 1) * 32;
}
if (w % 32 != 0) {
w = (w / 32 + 1) * 32;
}
cv::resize(In_Out_img, In_Out_img, cv::Size(w, h));
return scale;
}
// Creat the engine using only the API and not any parser.
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) {
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
// Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims4{ 1, 3, -1, -1 });
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("E:\\LearningCodes\\DBNET\\DBNet.pytorch\\tools\\DBNet.wts");
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
/* ------ Resnet18 backbone------ */
// Add convolution layer with 6 outputs and a 5x5 filter.
/* ------ Resnet18 backbone------ */
// Add convolution layer with 6 outputs and a 5x5 filter.
IConvolutionLayer* conv1 = network->addConvolution(*data, 64, DimsHW{ 7, 7 }, weightMap["backbone.conv1.weight"], emptywts);
assert(conv1);
conv1->setStride(DimsHW{ 2, 2 });
conv1->setPadding(DimsHW{ 3, 3 });
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "backbone.bn1", 1e-5);
// Add activation layer using the ReLU algorithm.
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
// Add max pooling layer with stride of 2x2 and kernel size of 2x2.
IPoolingLayer* pool1 = network->addPooling(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 });
assert(pool1);
pool1->setStride(DimsHW{ 2, 2 });
@ -63,171 +86,147 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
IActivationLayer* relu8 = basicBlock(network, weightMap, *relu7->getOutput(0), 256, 512, 2, "backbone.layer4.0.");
IActivationLayer* relu9 = basicBlock(network, weightMap, *relu8->getOutput(0), 512, 512, 1, "backbone.layer4.1."); //x5
/* ------- FPN neck ------- */
// net weight input,outch, ksize, s, g, std::string lname
// 1
auto p5 = convBnLeaky(network, weightMap, *relu9->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c5"); // k=1 s = 1 p = k/2=1/2=0
auto c4_1 = convBnLeaky(network, weightMap, *relu7->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c4");
/* ------- FPN neck ------- */
ILayer* p5 = convBnLeaky(network, weightMap, *relu9->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c5.conv", ".bn"); // k=1 s = 1 p = k/2=1/2=0
ILayer* c4_1 = convBnLeaky(network, weightMap, *relu7->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c4.conv", ".bn");
float *deval = reinterpret_cast<float*>(malloc(sizeof(float) * 64 * 2 * 2));
for (int i = 0; i < 64 * 2 * 2; i++) {
deval[i] = 1.0;
deval[i] = 1.0;
}
Weights deconvwts1{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p4_1 = network->addDeconvolutionNd(*p5->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts1, emptywts);
p4_1->setStrideNd(DimsHW{ 2, 2 });
p4_1->setNbGroups(64);
weightMap["deconv1"] = deconvwts1;
p4_1->setNbGroups(64);
weightMap["deconv1"] = deconvwts1;
auto p4_add = network->addElementWise(*p4_1->getOutput(0), *c4_1->getOutput(0), ElementWiseOperation::kSUM);
auto p4 = convBnLeaky(network, weightMap, *p4_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p4"); // smooth
// 2
auto c3_1 = convBnLeaky(network, weightMap, *relu5->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c3");
IElementWiseLayer* p4_add = network->addElementWise(*p4_1->getOutput(0), *c4_1->getOutput(0), ElementWiseOperation::kSUM);
ILayer* p4 = convBnLeaky(network, weightMap, *p4_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p4.conv", ".bn"); // smooth
ILayer* c3_1 = convBnLeaky(network, weightMap, *relu5->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c3.conv", ".bn");
Weights deconvwts2{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p3_1 = network->addDeconvolutionNd(*p4->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts2, emptywts);
p3_1->setStrideNd(DimsHW{ 2, 2 });
p3_1->setNbGroups(64);
Weights deconvwts2{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p3_1 = network->addDeconvolutionNd(*p4->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts2, emptywts);
p3_1->setStrideNd(DimsHW{ 2, 2 });
p3_1->setNbGroups(64);
auto p3_add = network->addElementWise(*p3_1->getOutput(0), *c3_1->getOutput(0), ElementWiseOperation::kSUM);
auto p3 = convBnLeaky(network, weightMap, *p3_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p3"); // smooth
// 3
auto c2_1 = convBnLeaky(network, weightMap, *relu3->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c2");
IElementWiseLayer* p3_add = network->addElementWise(*p3_1->getOutput(0), *c3_1->getOutput(0), ElementWiseOperation::kSUM);
ILayer* p3 = convBnLeaky(network, weightMap, *p3_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p3.conv", ".bn"); // smooth
ILayer* c2_1 = convBnLeaky(network, weightMap, *relu3->getOutput(0), 64, 1, 1, 1, "neck.reduce_conv_c2.conv", ".bn");
Weights deconvwts3{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p2_1 = network->addDeconvolutionNd(*p3->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts3, emptywts);
p2_1->setStrideNd(DimsHW{ 2, 2 });
p2_1->setNbGroups(64);
//Dims p2_1dim = p2_1->getOutput(0)->getDimensions();
auto p2_add = network->addElementWise(*p2_1->getOutput(0), *c2_1->getOutput(0), ElementWiseOperation::kSUM);
auto p2 = convBnLeaky(network, weightMap, *p2_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p2"); // smooth
Weights deconvwts3{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p2_1 = network->addDeconvolutionNd(*p3->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts3, emptywts);
p2_1->setStrideNd(DimsHW{ 2, 2 });
p2_1->setNbGroups(64);
IElementWiseLayer* p2_add = network->addElementWise(*p2_1->getOutput(0), *c2_1->getOutput(0), ElementWiseOperation::kSUM);
ILayer* p2 = convBnLeaky(network, weightMap, *p2_add->getOutput(0), 64, 3, 1, 1, "neck.smooth_p2.conv", ".bn"); // smooth
Weights deconvwts4{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p3_up_p2 = network->addDeconvolutionNd(*p3->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts4, emptywts);
p3_up_p2->setStrideNd(DimsHW{ 2, 2 });
p3_up_p2->setNbGroups(64);
// _upsample_cat
// p3--p2 (upx2 w p s=2 0 2)
Weights deconvwts4{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* p3_up_p2 = network->addDeconvolutionNd(*p3->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts4, emptywts);
p3_up_p2->setStrideNd(DimsHW{ 2, 2 });
p3_up_p2->setNbGroups(64);
float *deval2 = reinterpret_cast<float*>(malloc(sizeof(float) * 64 * 8 * 8));
for (int i = 0; i < 64 * 8 * 8; i++) {
deval2[i] = 1.0;
}
Weights deconvwts5{ DataType::kFLOAT, deval2, 64 * 8 * 8 };
IDeconvolutionLayer* p4_up_p2 = network->addDeconvolutionNd(*p4->getOutput(0), 64, DimsHW{ 8, 8 }, deconvwts5, emptywts);
p4_up_p2->setPadding(DimsHW{ 2, 2 });
p4_up_p2->setStrideNd(DimsHW{ 4, 4 });
p4_up_p2->setNbGroups(64);
weightMap["deconv2"] = deconvwts5;
// p4--p2(upx4 wps=824)
float *deval2 = reinterpret_cast<float*>(malloc(sizeof(float) * 64 * 8 * 8));
for (int i = 0; i < 64 * 8 * 8; i++) {
deval2[i] = 1.0;
}
Weights deconvwts5{ DataType::kFLOAT, deval2, 64 * 8 * 8 };
IDeconvolutionLayer* p4_up_p2 = network->addDeconvolutionNd(*p4->getOutput(0), 64, DimsHW{ 8, 8 }, deconvwts5, emptywts);
p4_up_p2->setPadding(DimsHW{ 2, 2 });
p4_up_p2->setStrideNd(DimsHW{ 4, 4 });
p4_up_p2->setNbGroups(64);
weightMap["deconv2"] = deconvwts5;
Weights deconvwts6{ DataType::kFLOAT, deval2, 64 * 8 * 8 };
IDeconvolutionLayer* p5_up_p2 = network->addDeconvolutionNd(*p5->getOutput(0), 64, DimsHW{ 8, 8 }, deconvwts6, emptywts);
p5_up_p2->setStrideNd(DimsHW{ 8, 8 });
p5_up_p2->setNbGroups(64);
// p5--p2(upx8) wps =808
Weights deconvwts6{ DataType::kFLOAT, deval2, 64 * 8 * 8 };
IDeconvolutionLayer* p5_up_p2 = network->addDeconvolutionNd(*p5->getOutput(0), 64, DimsHW{ 8, 8 }, deconvwts6, emptywts);
p5_up_p2->setStrideNd(DimsHW{ 8, 8 });
p5_up_p2->setNbGroups(64);
// torch.cat([p2, p3, p4, p5], dim=1)
ITensor* inputTensors[] = { p2->getOutput(0), p3_up_p2->getOutput(0), p4_up_p2->getOutput(0), p5_up_p2->getOutput(0) };
IConcatenationLayer* neck_cat = network->addConcatenation(inputTensors, 4);
// torch.cat([p2, p3, p4, p5], dim=1)
//Dims p2dim = p2->getOutput(0)->getDimensions();
//Dims p3dim = p3_up_p2->getOutput(0)->getDimensions();
//Dims p4dim = p4_up_p2->getOutput(0)->getDimensions();
//Dims p5dim = p5_up_p2->getOutput(0)->getDimensions();
ILayer* neck_out = convBnLeaky(network, weightMap, *neck_cat->getOutput(0), 256, 3, 1, 1, "neck.conv.0", ".1"); // smooth
assert(neck_out);
ILayer* binarize1 = convBnLeaky(network, weightMap, *neck_out->getOutput(0), 64, 3, 1, 1, "head.binarize.0", ".1"); //
Weights deconvwts7{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* binarizeup = network->addDeconvolutionNd(*binarize1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts7, emptywts);
binarizeup->setStrideNd(DimsHW{ 2, 2 });
binarizeup->setNbGroups(64);
IScaleLayer* binarizebn1 = addBatchNorm2d(network, weightMap, *binarizeup->getOutput(0), "head.binarize.4", 1e-5);
IActivationLayer* binarizerelu1 = network->addActivation(*binarizebn1->getOutput(0), ActivationType::kRELU);
assert(binarizerelu1);
ITensor* inputTensors[] = { p2->getOutput(0), p3_up_p2->getOutput(0), p4_up_p2->getOutput(0), p5_up_p2->getOutput(0) };
auto neck_cat = network->addConcatenation(inputTensors, 4);
//Dims neck_catdim = neck_cat->getOutput(0)->getDimensions();
Weights deconvwts8{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* binarizeup2 = network->addDeconvolutionNd(*binarizerelu1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts8, emptywts);
binarizeup2->setStrideNd(DimsHW{ 2, 2 });
binarizeup2->setNbGroups(64);
ILayer* neck_out = convBnLeaky2(network, weightMap, *neck_cat->getOutput(0), 256, 3, 1, 1, "neck.conv"); // smooth
assert(neck_out);
//Dims neck_outdim = neck_out->getOutput(0)->getDimensions();
/* ------- head ------- */
// shrink_maps = self.binarize(x)
// net weight input,outch, ksize, s, g, std::string lname
auto binarize1 = convBnLeaky2(network, weightMap, *neck_out->getOutput(0), 64, 3, 1, 1, "head.binarize"); //
Weights deconvwts7{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* binarizeup = network->addDeconvolutionNd(*binarize1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts7, emptywts);
binarizeup->setStrideNd(DimsHW{ 2, 2 });
binarizeup->setNbGroups(64);
IScaleLayer* binarizebn1 = addBatchNorm2d(network, weightMap, *binarizeup->getOutput(0), "head.binarize.4", 1e-5);
IActivationLayer* binarizerelu1 = network->addActivation(*binarizebn1->getOutput(0), ActivationType::kRELU);
assert(binarizerelu1);
IConvolutionLayer* binarize3 = network->addConvolution(*binarizeup2->getOutput(0), 1, DimsHW{ 3, 3 }, weightMap["head.binarize.7.weight"], weightMap["head.binarize.7.bias"]);
assert(binarize3);
binarize3->setStride(DimsHW{ 1, 1 });
binarize3->setPadding(DimsHW{ 1, 1 });
IActivationLayer* binarize4 = network->addActivation(*binarize3->getOutput(0), ActivationType::kSIGMOID);
assert(binarize4);
Weights deconvwts8{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* binarizeup2 = network->addDeconvolutionNd(*binarizerelu1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts8, emptywts);
binarizeup2->setStrideNd(DimsHW{ 2, 2 });
binarizeup2->setNbGroups(64);
IConvolutionLayer* binarize3 = network->addConvolution(*binarizeup2->getOutput(0), 1, DimsHW{ 3, 3 }, weightMap["head.binarize.7.weight"], weightMap["head.binarize.7.bias"]);
assert(binarize3);
binarize3->setStride(DimsHW{ 1, 1 });
binarize3->setPadding(DimsHW{ 1, 1 });
IActivationLayer* binarize4 = network->addActivation(*binarize3->getOutput(0), ActivationType::kSIGMOID);
assert(binarize4);
//threshold_maps = self.thresh(x)
ILayer* thresh1 = convBnLeaky(network, weightMap, *neck_out->getOutput(0), 64, 3, 1, 1, "head.thresh.0", ".1", false); //
Weights deconvwts9{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* threshup = network->addDeconvolutionNd(*thresh1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts9, emptywts);
threshup->setStrideNd(DimsHW{ 2, 2 });
threshup->setNbGroups(64);
IConvolutionLayer* thresh2 = network->addConvolution(*threshup->getOutput(0), 64, DimsHW{ 3, 3 }, weightMap["head.thresh.3.1.weight"], weightMap["head.thresh.3.1.bias"]);
assert(thresh2);
thresh2->setStride(DimsHW{ 1, 1 });
thresh2->setPadding(DimsHW{ 1, 1 });
//threshold_maps = self.thresh(x)
auto thresh1 = convBnLeaky2(network, weightMap, *neck_out->getOutput(0), 64, 3, 1, 1, "head.thresh", false); //
Weights deconvwts9{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* threshup = network->addDeconvolutionNd(*thresh1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts9, emptywts);
threshup->setStrideNd(DimsHW{ 2, 2 });
threshup->setNbGroups(64);
IConvolutionLayer* thresh2 = network->addConvolution(*threshup->getOutput(0), 64, DimsHW{ 3, 3 }, weightMap["head.thresh.3.1.weight"], weightMap["head.thresh.3.1.bias"]);
assert(thresh2);
thresh2->setStride(DimsHW{ 1, 1 });
thresh2->setPadding(DimsHW{ 1, 1 });
IScaleLayer* threshbn1 = addBatchNorm2d(network, weightMap, *thresh2->getOutput(0), "head.thresh.4", 1e-5);
IActivationLayer* threshrelu1 = network->addActivation(*threshbn1->getOutput(0), ActivationType::kRELU);
assert(threshrelu1);
IScaleLayer* threshbn1 = addBatchNorm2d(network, weightMap, *thresh2->getOutput(0), "head.thresh.4", 1e-5);
IActivationLayer* threshrelu1 = network->addActivation(*threshbn1->getOutput(0), ActivationType::kRELU);
assert(threshrelu1);
Weights deconvwts10{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* threshup2 = network->addDeconvolutionNd(*threshrelu1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts10, emptywts);
threshup2->setStrideNd(DimsHW{ 2, 2 });
threshup2->setNbGroups(64);
IConvolutionLayer* thresh3 = network->addConvolution(*threshup2->getOutput(0), 1, DimsHW{ 3, 3 }, weightMap["head.thresh.6.1.weight"], weightMap["head.thresh.6.1.bias"]);
assert(thresh3);
thresh3->setStride(DimsHW{ 1, 1 });
thresh3->setPadding(DimsHW{ 1, 1 });
IActivationLayer* thresh4 = network->addActivation(*thresh3->getOutput(0), ActivationType::kSIGMOID);
assert(thresh4);
Weights deconvwts10{ DataType::kFLOAT, deval, 64 * 2 * 2 };
IDeconvolutionLayer* threshup2 = network->addDeconvolutionNd(*threshrelu1->getOutput(0), 64, DimsHW{ 2, 2 }, deconvwts10, emptywts);
threshup2->setStrideNd(DimsHW{ 2, 2 });
threshup2->setNbGroups(64);
IConvolutionLayer* thresh3 = network->addConvolution(*threshup2->getOutput(0), 1, DimsHW{ 3, 3 }, weightMap["head.thresh.6.1.weight"], weightMap["head.thresh.6.1.bias"]);
assert(thresh3);
thresh3->setStride(DimsHW{ 1, 1 });
thresh3->setPadding(DimsHW{ 1, 1 });
IActivationLayer* thresh4 = network->addActivation(*thresh3->getOutput(0), ActivationType::kSIGMOID);
assert(thresh4);
ITensor* inputTensors2[] = { binarize4->getOutput(0), thresh4->getOutput(0) };
IConcatenationLayer* head_out = network->addConcatenation(inputTensors2, 2);
//y = torch.cat((shrink_maps, threshold_maps), dim=1)
// binarize4 thresh4
//Dims binarize4dim = binarize4->getOutput(0)->getDimensions();
//Dims thresh4dim = thresh4->getOutput(0)->getDimensions();
// y = F.interpolate(y, size=(H, W))
head_out->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*head_out->getOutput(0));
ITensor* inputTensors2[] = { binarize4->getOutput(0), thresh4->getOutput(0)};
auto head_out = network->addConcatenation(inputTensors2, 2);
IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kMIN, Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE));
profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kOPT, Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W));
profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kMAX, Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE));
config->addOptimizationProfile(profile);
// y = F.interpolate(y, size=(H, W)) # 使用最近邻训练的可以用TRTAPI实现
// 最后大小为图片大小
head_out->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*head_out->getOutput(0));
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
// Build engine
builder->setMaxBatchSize(maxBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
config->setFlag(BuilderFlag::kFP16);
#endif
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
std::cout << "Building engine, please wait for a while..." << std::endl;
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "Build engine successfully!" << std::endl;
// Don't need the network any more
network->destroy();
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*)(mem.second.values));
}
// Release host memory
for (auto& mem : weightMap) {
free((void*)(mem.second.values));
}
return engine;
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
@ -236,8 +235,8 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
//ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
//ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
@ -248,7 +247,7 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) {
builder->destroy();
}
void doInference(IExecutionContext& context, float* input, float* output, int batchSize) {
void doInference(IExecutionContext& context, float* input, float* output, int h_scale, int w_scale) {
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
@ -260,19 +259,20 @@ void doInference(IExecutionContext& context, float* input, float* output, int ba
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
context.setBindingDimensions(inputIndex, Dims4(1, 3, h_scale, w_scale));
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
CHECK(cudaMalloc(&buffers[inputIndex], 3 * h_scale * w_scale * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], 2 * h_scale * w_scale * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, 3 * h_scale * w_scale * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], h_scale * w_scale * 2 * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
@ -281,149 +281,130 @@ void doInference(IExecutionContext& context, float* input, float* output, int ba
CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv)
{
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{ nullptr };
size_t size{ 0 };
int main(int argc, char** argv) {
cudaSetDevice(DEVICE);
// create a model using the API directly and serialize it to a stream
char *trtModelStream{ nullptr };
size_t size{ 0 };
if (argc == 2 && std::string(argv[1]) == "-s")
{
IHostMemory* modelStream{ nullptr };
APIToModel(BATCH_SIZE, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("DBNet.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
if (argc == 2 && std::string(argv[1]) == "-s") {
IHostMemory* modelStream{ nullptr };
APIToModel(1, &modelStream);
assert(modelStream != nullptr);
std::ofstream p("DBNet.engine", std::ios::binary);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
else if (argc == 3 && std::string(argv[1]) == "-d")
{
std::ifstream file("DBNet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
modelStream->destroy();
return 0;
}
else if (argc == 3 && std::string(argv[1]) == "-d") {
std::ifstream file("DBNet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
else
{
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./debnet -s // serialize model to plan file" << std::endl;
std::cerr << "./debnet -d ../samples // deserialize plan file and run inference" << std::endl;
return -1;
}
else {
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./debnet -s // serialize model to plan file" << std::endl;
std::cerr << "./debnet -d ../samples // deserialize plan file and run inference" << std::endl;
return -1;
}
// prepare input data ---------------------------
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
std::vector<std::string> file_names;
if (read_files_in_dir(argv[2], file_names) < 0) {
std::cout << "read_files_in_dir failed." << std::endl;
return -1;
}
std::vector<float> mean_value{ 0.406, 0.456, 0.485 }; // BGR
std::vector<float> std_value{ 0.225, 0.224, 0.229 };
int fcount = 0;
for (auto f : file_names) {
fcount++;
std::cout << fcount << " " << f << std::endl;
cv::Mat pr_img = cv::imread(std::string(argv[2]) + "/" + f);
cv::Mat src_img = pr_img.clone();
if (pr_img.empty()) continue;
float scale = paddimg(pr_img, SHORT_INPUT);
std::cout << "letterbox shape: " << pr_img.cols << ", " << pr_img.rows << std::endl;
if (pr_img.cols < MIN_INPUT_SIZE || pr_img.rows < MIN_INPUT_SIZE) continue;
float* data = new float[3 * pr_img.rows * pr_img.cols];
int i = 0;
for (int row = 0; row < pr_img.rows; ++row) {
uchar* uc_pixel = pr_img.data + row * pr_img.step;
for (int col = 0; col < pr_img.cols; ++col) {
data[i] = (uc_pixel[2] / 255.0 - mean_value[2]) / std_value[2];
data[i + pr_img.rows * pr_img.cols] = (uc_pixel[1] / 255.0 - mean_value[1]) / std_value[1];
data[i + 2 * pr_img.rows * pr_img.cols] = (uc_pixel[0] / 255.0 - mean_value[0]) / std_value[0];
uc_pixel += 3;
++i;
}
}
// prepare input data ---------------------------
static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
//for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
// data[i] = 1.0;
static float prob[BATCH_SIZE * OUTPUT_SIZE];
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
float* prob = new float[pr_img.rows *pr_img.cols * 2];
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, pr_img.rows, pr_img.cols);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
std::vector<std::string> file_names;
if (read_files_in_dir(argv[2], file_names) < 0) {
std::cout << "read_files_in_dir failed." << std::endl;
return -1;
// prob 为 2* 640*640 拿出第一个
cv::Mat map = cv::Mat::zeros(cv::Size(pr_img.cols, pr_img.rows), CV_8UC1);
for (int h = 0; h < pr_img.rows; ++h) {
uchar *ptr = map.ptr(h);
for (int w = 0; w < pr_img.cols; ++w) {
ptr[w] = (prob[h * pr_img.cols + w] > 0.3) ? 255 : 0;
}
}
/*
std::vector<float> mean_value{0.406, 0.456, 0.485};
std::vector<float> std_value{0.225, 0.224, 0.229};
cv::Mat src, dst;
std::vector<cv::Mat> bgrChannels(3);
cv::split(src, bgrChannels);
for (auto i = 0; i < bgrChannels.size(); i++)
{
bgrChannels[i].convertTo(bgrChannels[i], CV_32FC1, 1.0 / std_value[i], (0.0 - mean_value[i]) / std_value[i]);
// 提取最小外接矩形
std::vector<std::vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarcy;
cv::findContours(map, contours, hierarcy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);
std::vector<cv::Rect> boundRect(contours.size());
std::vector<cv::RotatedRect> box(contours.size());
cv::Point2f rect[4];
for (int i = 0; i < contours.size(); i++) {
box[i] = cv::minAreaRect(cv::Mat(contours[i]));
//boundRect[i] = cv::boundingRect(cv::Mat(contours[i]));
//绘制外接矩形和 最小外接矩形（for循环）
//cv::rectangle(img, cv::Point(boundRect[i].x, boundRect[i].y), cv::Point(boundRect[i].x + boundRect[i].width, boundRect[i].y + boundRect[i].height), cv::Scalar(0, 255, 0), 2, 8);
cv::RotatedRect expandbox = expandBox(box[i], EXPANDRATIO);
expandbox.points(rect);//把最小外接矩形四个端点复制给rect数组
for (int j = 0; j < 4; j++) {
cv::Point2f p1, p2;
p1.x = round(rect[j].x / pr_img.cols * src_img.cols);
p1.y = round(rect[j].y / pr_img.rows * src_img.rows);
p2.x = round(rect[(j + 1) % 4].x / pr_img.cols * src_img.cols);
p2.y = round(rect[(j + 1) % 4].y / pr_img.rows * src_img.rows);
cv::line(src_img, p1, p2, cv::Scalar(0, 0, 255), 2, 8);
}
}
cv::meger(bgrChannels, dst);
*/
std::vector<float> mean_value{ 0.406, 0.456, 0.485 }; // BGR
std::vector<float> std_value{ 0.225, 0.224, 0.229 };
int fcount = 0;
for (int f = 0; f < (int)file_names.size(); f++) {
fcount++;
if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue;
for (int b = 0; b < fcount; b++) {
//cv::Mat img = cv::imread(file_names[f - fcount + 1 + b]);
cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]);
if (img.empty()) continue;
cv::Mat pr_img; // letterbox BGR to RGB
cv::resize(img, pr_img, cv::Size(INPUT_W, INPUT_H), cv::INTER_LINEAR);
int i = 0;
for (int row = 0; row < INPUT_H; ++row) {
uchar* uc_pixel = pr_img.data + row * pr_img.step;
for (int col = 0; col < INPUT_W; ++col) {
data[b * 3 * INPUT_H * INPUT_W + i] = (uc_pixel[2]/255.0 - mean_value[2]) / std_value[2];
data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (uc_pixel[1]/255.0 - mean_value[1]) / std_value[1];
data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (uc_pixel[0]/255.0 - mean_value[0]) / std_value[0];
uc_pixel += 3;
++i;
}
}
}
cv::imwrite("_" + f, src_img);
//cv::waitKey(0);
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, BATCH_SIZE);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
delete prob;
delete data;
}
// prob 为 2* 640*640 拿出第一个
cv::Mat map = cv::Mat::zeros(cv::Size(640, 640), CV_8UC1);
for (int b = 0; b < fcount; b++)
{
cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]);
cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), cv::INTER_LINEAR);
for (int h = 0; h < INPUT_H; ++h)
{
uchar *ptr = map.ptr(h);
for (int w = 0; w < INPUT_W; ++w)
{
ptr[w] = (prob[b*OUTPUT_SIZE + h*INPUT_W + w] > 0.3) ? 255 : 0;
}
}
// 提取最小外接矩形
std::vector<std::vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarcy;
cv::findContours(map, contours, hierarcy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);
std::vector<cv::Rect> boundRect(contours.size());
std::vector<cv::RotatedRect> box(contours.size());
cv::Point2f rect[4];
for (int i = 0; i < contours.size(); i++)
{
box[i] = cv::minAreaRect(cv::Mat(contours[i]));
//boundRect[i] = cv::boundingRect(cv::Mat(contours[i]));
//绘制外接矩形和 最小外接矩形（for循环）
//cv::rectangle(img, cv::Point(boundRect[i].x, boundRect[i].y), cv::Point(boundRect[i].x + boundRect[i].width, boundRect[i].y + boundRect[i].height), cv::Scalar(0, 255, 0), 2, 8);
cv::RotatedRect expandbox = expandBox(box[i], EXPANDRATIO);
expandbox.points(rect);//把最小外接矩形四个端点复制给rect数组
for (int j = 0; j < 4; j++)
{
cv::line(img, rect[j], rect[(j + 1) % 4], cv::Scalar(0, 0, 255), 2, 8);
}
}
cv::imshow("result", img);
cv::waitKey(0);
}
return 0;
}
return 0;
}