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

fix network bug and rewrite post-processing pse algorithm (#367)

Browse Source 
* create psenet

create psenet with weight from tensorflow

* delete some useless code

* repalce tab with 4 blanks

* fix network bug, rewrite post-processing pse algorithm

* update readme

* update readme
This commit is contained in:
weiwei zhou 2021-01-22 19:44:29 +08:00 committed by GitHub
parent 556665cca1
commit 285ec22e5f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 812 additions and 853 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
psenet/README.md
View File

@ -1,7 +1,7 @@
# PSENet
**preprocessing + inference + postprocessing = 30ms** with fp32 on Tesla P40.
The Tensorflow implementation is [tensorflow_PSENet](https://github.com/liuheng92/tensorflow_PSENet).
The original Tensorflow implementation is [tensorflow_PSENet](https://github.com/liuheng92/tensorflow_PSENet). A TensorRT Python api implementation is [TensorRT-Python-PSENet](https://github.com/upczww/TensorRT-Python-PSENet).
## Key Features
- Generating `.wts` from `Tensorflow`.
@ -9,9 +9,8 @@ The Tensorflow implementation is [tensorflow_PSENet](https://github.com/liuheng9
- Object-Oriented Programming.
- Practice with C++ 11.
<p align="center">
<img src="https://user-images.githubusercontent.com/15235574/101176279-3868b780-3681-11eb-9ca3-7425b41fe895.jpg">
</p>
## How to Run
@ -42,14 +41,12 @@ The Tensorflow implementation is [tensorflow_PSENet](https://github.com/liuheng9
cp ../psenet.wts ./
cp ../test.jpg ./
./psenet -s // serialize model to plan file
./psenet -d // deserialize plan file and run inference"
./psenet -d // deserialize plan file and run inference
```
## Known Issues
1. The output of network is not completely the same as the tf's due to the difference between tensorrt's `addResize` and `tf.image.resize`, I will figure it out.
None
## Todo
* use `ExponentialMovingAverage` weight.
* faster preporcess and postprocess.

60
psenet/gen_tf_wts.py
View File

@ -1,31 +1,31 @@
from sys import prefix
import tensorflow as tf
from tensorflow.python import pywrap_tensorflow
import numpy as np
import struct
model_dir = "model"
ckpt = tf.train.get_checkpoint_state(model_dir)
ckpt_path = ckpt.model_checkpoint_path
reader = pywrap_tensorflow.NewCheckpointReader(ckpt_path)
param_dict = reader.get_variable_to_shape_map()
f = open(r"psenet.wts", "w")
keys = param_dict.keys()
f.write("{}\n".format(len(keys)))
for key in keys:
weight = reader.get_tensor(key)
print(key, weight.shape)
if len(weight.shape) == 4:
weight = np.transpose(weight, (3, 2, 0, 1))
print(weight.shape)
weight = np.reshape(weight, -1)
f.write("{} {} ".format(key, len(weight)))
for w in weight:
f.write(" ")
f.write(struct.pack(">f", float(w)).hex())
from sys import prefix
import tensorflow as tf
from tensorflow.python import pywrap_tensorflow
import numpy as np
import struct
model_dir = "model"
ckpt = tf.train.get_checkpoint_state(model_dir)
ckpt_path = ckpt.model_checkpoint_path
reader = pywrap_tensorflow.NewCheckpointReader(ckpt_path)
param_dict = reader.get_variable_to_shape_map()
f = open(r"psenet.wts", "w")
keys = param_dict.keys()
f.write("{}\n".format(len(keys)))
for key in keys:
weight = reader.get_tensor(key)
print(key, weight.shape)
if len(weight.shape) == 4:
weight = np.transpose(weight, (3, 2, 0, 1))
print(weight.shape)
weight = np.reshape(weight, -1)
f.write("{} {} ".format(key, len(weight)))
for w in weight:
f.write(" ")
f.write(struct.pack(">f", float(w)).hex())
f.write("\n")

244
psenet/layers.cpp
View File

@ -1,136 +1,110 @@
#include "layers.h"
IScaleLayer *addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, std::string lname, float eps)
{
float *gamma = (float *)weightMap[lname + "gamma"].values; // scale
float *beta = (float *)weightMap[lname + "beta"].values; // offset
float *mean = (float *)weightMap[lname + "moving_mean"].values;
float *var = (float *)weightMap[lname + "moving_variance"].values;
int len = weightMap[lname + "moving_variance"].count;
float *scval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
for (auto i = 0; i < len; i++)
{
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{DataType::kFLOAT, scval, len};
float *shval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
for (auto i = 0; i < len; i++)
{
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{DataType::kFLOAT, shval, len};
float *pval = reinterpret_cast<float *>(malloc(sizeof(float) * len));
for (auto i = 0; i < len; i++)
{
pval[i] = 1.0;
}
Weights power{DataType::kFLOAT, pval, len};
IScaleLayer *scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IActivationLayer *bottleneck(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int ch, int stride, std::string lname, int branch_type)
{
Weights emptywts{DataType::kFLOAT, nullptr, 0};
IConvolutionLayer *conv1 = network->addConvolution(input, ch, DimsHW{1, 1}, weightMap[lname + "conv1/weights"], emptywts);
assert(conv1);
Dims conv1_shape = conv1->getOutput(0)->getDimensions();
IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "conv1/BatchNorm/", 1e-5);
assert(bn1);
Dims bn1_shape = bn1->getOutput(0)->getDimensions();
IActivationLayer *relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
Dims relu1_shape = relu1->getOutput(0)->getDimensions();
IConvolutionLayer *conv2 = network->addConvolution(*relu1->getOutput(0), ch, DimsHW{3, 3}, weightMap[lname + "conv2/weights"], emptywts);
assert(conv2);
conv2->setStride(DimsHW{stride, stride});
conv2->setPadding(DimsHW{1, 1});
Dims conv2_shape = conv2->getOutput(0)->getDimensions();
IScaleLayer *bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "conv2/BatchNorm/", 1e-5);
assert(bn2);
Dims bn2_shape = bn2->getOutput(0)->getDimensions();
IActivationLayer *relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU);
assert(relu2);
Dims relu2_shape = relu2->getOutput(0)->getDimensions();
IConvolutionLayer *conv3 = network->addConvolution(*relu2->getOutput(0), ch * 4, DimsHW{1, 1}, weightMap[lname + "conv3/weights"], emptywts);
assert(conv3);
Dims conv3_shape = conv3->getOutput(0)->getDimensions();
IScaleLayer *bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + "conv3/BatchNorm/", 1e-5);
assert(bn3);
IElementWiseLayer *ew1;
Dims ew1_shape;
// branch_type 0:shortcut,1:conv+bn+shortcut,2:maxpool+shortcut
if (branch_type == 0)
{
ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
ew1_shape = ew1->getOutput(0)->getDimensions();
assert(ew1);
}
else if (branch_type == 1)
{
IConvolutionLayer *conv4 = network->addConvolution(input, ch * 4, DimsHW{1, 1}, weightMap[lname + "shortcut/weights"], emptywts);
assert(conv4);
conv4->setStride(DimsHW{stride, stride});
IScaleLayer *bn4 = addBatchNorm2d(network, weightMap, *conv4->getOutput(0), lname + "shortcut/BatchNorm/", 1e-5);
assert(bn4);
ew1 = network->addElementWise(*bn4->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
ew1_shape = ew1->getOutput(0)->getDimensions();
assert(ew1);
}
else
{
IPoolingLayer *pool = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{1, 1});
assert(pool);
pool->setStrideNd(DimsHW{2, 2});
ew1 = network->addElementWise(*pool->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
ew1_shape = ew1->getOutput(0)->getDimensions();
assert(ew1);
}
IActivationLayer *relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU);
Dims relu3_shape = relu3->getOutput(0)->getDimensions();
assert(relu3);
return relu3;
}
IActivationLayer *ConvRelu(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int outch, int kernel, int stride, std::string lname)
{
IConvolutionLayer *conv = network->addConvolution(input, 256, DimsHW{kernel, kernel}, weightMap[lname + "weights"], weightMap[lname + "biases"]);
assert(conv);
conv->setStride(DimsHW{stride, stride});
if (kernel == 3 || stride == 2)
{
conv->setPadding(DimsHW{1, 1});
}
IActivationLayer *ac = network->addActivation(*conv->getOutput(0), ActivationType::kRELU);
assert(ac);
return ac;
#include "layers.h"
IScaleLayer* addBatchNorm2d(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps)
{
float* gamma = (float*)weightMap[lname + "gamma"].values; // scale
float* beta = (float*)weightMap[lname + "beta"].values; // offset
float* mean = (float*)weightMap[lname + "moving_mean"].values;
float* var = (float*)weightMap[lname + "moving_variance"].values;
int len = weightMap[lname + "moving_variance"].count;
float* scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (auto i = 0; i < len; i++)
{
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
Weights scale{ DataType::kFLOAT, scval, len };
float* shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (auto i = 0; i < len; i++)
{
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
Weights shift{ DataType::kFLOAT, shval, len };
float* pval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for (auto i = 0; i < len; i++)
{
pval[i] = 1.0;
}
Weights power{ DataType::kFLOAT, pval, len };
IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
assert(scale_1);
return scale_1;
}
IActivationLayer* bottleneck(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int ch, int stride, std::string lname, int branch_type)
{
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
IConvolutionLayer* conv1 = network->addConvolutionNd(input, ch, DimsHW{ 1, 1 }, weightMap[lname + "conv1/weights"], emptywts);
assert(conv1);
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + "conv1/BatchNorm/", 1e-5);
assert(bn1);
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
IConvolutionLayer* conv2 = network->addConvolutionNd(*relu1->getOutput(0), ch, DimsHW{ 3, 3 }, weightMap[lname + "conv2/weights"], emptywts);
conv2->setStrideNd(DimsHW{ stride, stride });
conv2->setPaddingNd(DimsHW{ 1, 1 });
assert(conv2);
IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + "conv2/BatchNorm/", 1e-5);
assert(bn2);
IActivationLayer* relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU);
assert(relu2);
IConvolutionLayer* conv3 = network->addConvolutionNd(*relu2->getOutput(0), ch * 4, DimsHW{ 1, 1 }, weightMap[lname + "conv3/weights"], emptywts);
assert(conv3);
IScaleLayer* bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + "conv3/BatchNorm/", 1e-5);
assert(bn3);
IElementWiseLayer* ew1;
// branch_type 0:shortcut,1:conv+bn+shortcut,2:maxpool+shortcut
if (branch_type == 0)
{
ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
}
else if (branch_type == 1)
{
IConvolutionLayer* conv4 = network->addConvolutionNd(input, ch * 4, DimsHW{ 1, 1 }, weightMap[lname + "shortcut/weights"], emptywts);
conv4->setStrideNd(DimsHW{ stride, stride });
assert(conv4);
IScaleLayer* bn4 = addBatchNorm2d(network, weightMap, *conv4->getOutput(0), lname + "shortcut/BatchNorm/", 1e-5);
assert(bn4);
ew1 = network->addElementWise(*bn4->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
}
else
{
IPoolingLayer* pool = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{ 1, 1 });
pool->setStrideNd(DimsHW{ 2, 2 });
assert(pool);
ew1 = network->addElementWise(*pool->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM);
assert(ew1);
}
IActivationLayer* relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU);
assert(relu3);
return relu3;
}
IActivationLayer* addConvRelu(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int kernel, int stride, std::string lname)
{
IConvolutionLayer* conv = network->addConvolutionNd(input, 256, DimsHW{ kernel, kernel }, weightMap[lname + "weights"], weightMap[lname + "biases"]);
conv->setStrideNd(DimsHW{ stride, stride });
if (kernel == 3)
{
conv->setPaddingNd(DimsHW{ 1, 1 });
}
assert(conv);
IActivationLayer* ac = network->addActivation(*conv->getOutput(0), ActivationType::kRELU);
assert(ac);
return ac;
}

70
psenet/main.cpp
View File

@ -1,36 +1,34 @@
#include "psenet.h"
int main(int argc, char **argv)
{
PSENet psenet(1600, 0.9, 6, 4);
if (argc == 2 && std::string(argv[1]) == "-s")
{
std::cout << "Serializling Engine" << std::endl;
psenet.serializeEngine();
return 0;
}
else if (argc == 2 && std::string(argv[1]) == "-d")
{
psenet.init();
std::vector<std::string> files;
for (int i = 0; i < 10; i++)
{
files.emplace_back("test.jpg");
}
for (auto file : files)
{
std::cout << "Detect " << file << std::endl;
psenet.detect(file);
}
return 0;
}
else
{
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./psenet -s // serialize model to plan file" << std::endl;
std::cerr << "./psenet -d // deserialize plan file and run inference" << std::endl;
return -1;
}
}
#include "psenet.h"
int main(int argc, char** argv)
{
PSENet psenet(1200, 640, 0.90, 6, 4);
if (argc == 2 && std::string(argv[1]) == "-s")
{
std::cout << "Serializling Engine" << std::endl;
psenet.serializeEngine();
return 0;
}
else if (argc == 2 && std::string(argv[1]) == "-d")
{
psenet.init();
std::vector<std::string> files;
for (int i = 0; i < 10; i++)
files.emplace_back("test.jpg");
for (auto file : files)
{
std::cout << "Detect " << file << std::endl;
psenet.detect(file);
}
return 0;
}
else
{
std::cerr << "arguments not right!" << std::endl;
std::cerr << "./psenet -s // serialize model to plan file" << std::endl;
std::cerr << "./psenet -d // deserialize plan file and run inference" << std::endl;
return -1;
}
}

896
psenet/psenet.cpp
View File

@ -1,451 +1,445 @@
#include "psenet.h"
#define MAX_INPUT_SIZE 1200
#define MIN_INPUT_SIZE 128
#define OPT_INPUT_W 640
#define OPT_INPUT_H 640
PSENet::PSENet(int max_side_len, float threshold, int num_kernel, int stride) : max_side_len_(max_side_len),
post_threshold_(threshold),
num_kernels_(num_kernel),
stride_(stride)
{
}
PSENet::~PSENet()
{
}
// create the engine using only the API and not any parser.
ICudaEngine *PSENet::createEngine(IBuilder *builder, IBuilderConfig *config)
{
std::map<std::string, Weights> weightMap = loadWeights("./psenet.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
ITensor *data = network->addInput(input_name_, dt, Dims4{-1, 3, -1, -1});
assert(data);
IConvolutionLayer *conv1 = network->addConvolutionNd(*data, 64, DimsHW{7, 7}, weightMap["resnet_v1_50/conv1/weights"], emptywts);
conv1->setStrideNd(DimsHW{2, 2});
conv1->setPaddingNd(DimsHW{3, 3});
assert(conv1);
IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "resnet_v1_50/conv1/BatchNorm/", 1e-5);
assert(bn1);
IActivationLayer *relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
// C2
IPoolingLayer *pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
pool1->setStrideNd(DimsHW{2, 2});
pool1->setPaddingNd(DimsHW{1, 1});
assert(pool1);
IActivationLayer *x;
x = bottleneck(network, weightMap, *pool1->getOutput(0), 64, 1, "resnet_v1_50/block1/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 64, 1, "resnet_v1_50/block1/unit_2/bottleneck_v1/", 0);
// C3
IActivationLayer *block1 = bottleneck(network, weightMap, *x->getOutput(0), 64, 2, "resnet_v1_50/block1/unit_3/bottleneck_v1/", 2);
x = bottleneck(network, weightMap, *block1->getOutput(0), 128, 1, "resnet_v1_50/block2/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 128, 1, "resnet_v1_50/block2/unit_2/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 128, 1, "resnet_v1_50/block2/unit_3/bottleneck_v1/", 0);
// C4
IActivationLayer *block2 = bottleneck(network, weightMap, *x->getOutput(0), 128, 2, "resnet_v1_50/block2/unit_4/bottleneck_v1/", 2);
x = bottleneck(network, weightMap, *block2->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_2/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_3/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_4/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_5/bottleneck_v1/", 0);
IActivationLayer *block3 = bottleneck(network, weightMap, *x->getOutput(0), 256, 2, "resnet_v1_50/block3/unit_6/bottleneck_v1/", 2);
x = bottleneck(network, weightMap, *block3->getOutput(0), 512, 1, "resnet_v1_50/block4/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 512, 1, "resnet_v1_50/block4/unit_2/bottleneck_v1/", 0);
// C5
IActivationLayer *block4 = bottleneck(network, weightMap, *x->getOutput(0), 512, 1, "resnet_v1_50/block4/unit_3/bottleneck_v1/", 0);
IActivationLayer *build_p5_r1 = ConvRelu(network, weightMap, *block4->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P5/");
assert(build_p5_r1);
IActivationLayer *build_p4_r1 = ConvRelu(network, weightMap, *block2->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P4/reduce_dimension/");
assert(build_p4_r1);
IResizeLayer *bfp_layer4_resize = network->addResize(*build_p5_r1->getOutput(0));
auto build_p4_r1_shape = network->addShape(*build_p4_r1->getOutput(0))->getOutput(0);
bfp_layer4_resize->setInput(1, *build_p4_r1_shape);
bfp_layer4_resize->setResizeMode(ResizeMode::kNEAREST);
bfp_layer4_resize->setAlignCorners(false);
assert(bfp_layer4_resize);
IElementWiseLayer *bfp_add = network->addElementWise(*bfp_layer4_resize->getOutput(0), *build_p4_r1->getOutput(0), ElementWiseOperation::kSUM);
assert(bfp_add);
IActivationLayer *build_p4_r2 = ConvRelu(network, weightMap, *bfp_add->getOutput(0), 256, 3, 1, "build_feature_pyramid/build_P4/avoid_aliasing/");
assert(build_p4_r2);
IActivationLayer *build_p3_r1 = ConvRelu(network, weightMap, *block1->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P3/reduce_dimension/");
assert(build_p3_r1);
IResizeLayer *bfp_layer3_resize = network->addResize(*build_p4_r2->getOutput(0));
bfp_layer3_resize->setResizeMode(ResizeMode::kNEAREST);
auto build_p3_r1_shape = network->addShape(*build_p3_r1->getOutput(0))->getOutput(0);
bfp_layer3_resize->setInput(1, *build_p3_r1_shape);
bfp_layer3_resize->setAlignCorners(false);
assert(bfp_layer3_resize);
IElementWiseLayer *bfp_add1 = network->addElementWise(*bfp_layer3_resize->getOutput(0), *build_p3_r1->getOutput(0), ElementWiseOperation::kSUM);
assert(bfp_add1);
IActivationLayer *build_p3_r2 = ConvRelu(network, weightMap, *bfp_add1->getOutput(0), 256, 3, 1, "build_feature_pyramid/build_P3/avoid_aliasing/");
assert(build_p3_r2);
IActivationLayer *build_p2_r1 = ConvRelu(network, weightMap, *pool1->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P2/reduce_dimension/");
assert(build_p2_r1);
IResizeLayer *bfp_layer2_resize = network->addResize(*build_p3_r2->getOutput(0));
bfp_layer2_resize->setResizeMode(ResizeMode::kNEAREST);
auto build_p2_r1_shape = network->addShape(*build_p2_r1->getOutput(0))->getOutput(0);
bfp_layer2_resize->setInput(1, *build_p2_r1_shape);
bfp_layer2_resize->setAlignCorners(false);
assert(bfp_layer2_resize);
IElementWiseLayer *bfp_add2 = network->addElementWise(*bfp_layer2_resize->getOutput(0), *build_p2_r1->getOutput(0), ElementWiseOperation::kSUM);
assert(bfp_add2);
// P2
IActivationLayer *build_p2_r2 = ConvRelu(network, weightMap, *bfp_add2->getOutput(0), 256, 3, 1, "build_feature_pyramid/build_P2/avoid_aliasing/");
assert(build_p2_r2);
auto build_p2_r2_shape = network->addShape(*build_p2_r2->getOutput(0))->getOutput(0);
// P3 x2
IResizeLayer *layer1_resize = network->addResize(*build_p3_r2->getOutput(0));
layer1_resize->setResizeMode(ResizeMode::kLINEAR);
layer1_resize->setInput(1, *build_p2_r2_shape);
layer1_resize->setAlignCorners(true);
assert(layer1_resize);
// P4 x4
IResizeLayer *layer2_resize = network->addResize(*build_p4_r2->getOutput(0));
layer2_resize->setResizeMode(ResizeMode::kLINEAR);
layer2_resize->setInput(1, *build_p2_r2_shape);
layer2_resize->setAlignCorners(true);
assert(layer2_resize);
// P5 x8
IResizeLayer *layer3_resize = network->addResize(*build_p5_r1->getOutput(0));
layer3_resize->setResizeMode(ResizeMode::kLINEAR);
layer3_resize->setInput(1, *build_p2_r2_shape);
layer3_resize->setAlignCorners(true);
assert(layer3_resize);
// C(P5,P4,P3,P2)
ITensor *inputTensors[] = {layer3_resize->getOutput(0), layer2_resize->getOutput(0), layer1_resize->getOutput(0), build_p2_r2->getOutput(0)};
IConcatenationLayer *concat = network->addConcatenation(inputTensors, 4);
assert(concat);
IConvolutionLayer *feature_result_conv = network->addConvolutionNd(*concat->getOutput(0), 256, DimsHW{3, 3}, weightMap["feature_results/Conv/weights"], emptywts);
feature_result_conv->setPaddingNd(DimsHW{1, 1});
assert(feature_result_conv);
IScaleLayer *feature_result_bn = addBatchNorm2d(network, weightMap, *feature_result_conv->getOutput(0), "feature_results/Conv/BatchNorm/", 1e-5);
assert(feature_result_bn);
IActivationLayer *feature_result_relu = network->addActivation(*feature_result_bn->getOutput(0), ActivationType::kRELU);
assert(feature_result_relu);
IConvolutionLayer *feature_result_conv_1 = network->addConvolutionNd(*feature_result_relu->getOutput(0), 6, DimsHW{1, 1}, weightMap["feature_results/Conv_1/weights"], weightMap["feature_results/Conv_1/biases"]);
assert(feature_result_conv_1);
IActivationLayer *sigmoid = network->addActivation(*feature_result_conv_1->getOutput(0), ActivationType::kSIGMOID);
assert(sigmoid);
sigmoid->getOutput(0)->setName(output_name_);
std::cout << "Set name out" << std::endl;
network->markOutput(*sigmoid->getOutput(0));
// Set profile
IOptimizationProfile *profile = builder->createOptimizationProfile();
profile->setDimensions(input_name_, OptProfileSelector::kMIN, Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE));
profile->setDimensions(input_name_, OptProfileSelector::kOPT, Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W));
profile->setDimensions(input_name_, OptProfileSelector::kMAX, Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE));
config->addOptimizationProfile(profile);
// Build engine
config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
;
std::cout << "Build out" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto &mem : weightMap)
{
free((void *)(mem.second.values));
}
return engine;
}
void PSENet::serializeEngine()
{
// Create builder
IBuilder *builder = createInferBuilder(gLogger);
IBuilderConfig *config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine *engine = createEngine(builder, config);
assert(engine != nullptr);
// Serialize the engine
IHostMemory *modelStream{nullptr};
modelStream = engine->serialize();
assert(modelStream != nullptr);
std::ofstream p("./psenet.engine", std::ios::binary | std::ios::out);
if (!p)
{
std::cerr << "Could not open plan output file" << std::endl;
return;
}
p.write(reinterpret_cast<const char *>(modelStream->data()), modelStream->size());
return;
}
void PSENet::deserializeEngine()
{
std::ifstream file("./psenet.engine", std::ios::binary | std::ios::in);
if (file.good())
{
file.seekg(0, file.end);
size_t size = file.tellg();
file.seekg(0, file.beg);
char *trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
mCudaEngine = std::shared_ptr<nvinfer1::ICudaEngine>(mRuntime->deserializeCudaEngine(trtModelStream, size), InferDeleter());
assert(mCudaEngine != nullptr);
}
}
void PSENet::inferenceOnce(IExecutionContext &context, float *input, float *output, int input_h, int input_w)
{
const ICudaEngine &engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void *buffers[2];
// 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(input_name_);
const int outputIndex = engine.getBindingIndex(output_name_);
context.setBindingDimensions(inputIndex, Dims4(1, 3, input_h, input_w));
int input_size = 3 * input_h * input_w * sizeof(float);
int output_size = input_h * input_w * 6 / 16 * sizeof(float);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], input_size));
CHECK(cudaMalloc(&buffers[outputIndex], output_size));
// 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, input_size, cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size, cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
void PSENet::init()
{
mRuntime = std::shared_ptr<nvinfer1::IRuntime>(createInferRuntime(gLogger), InferDeleter());
assert(mRuntime != nullptr);
std::cout << "Deserialize Engine" << std::endl;
deserializeEngine();
mContext = std::shared_ptr<nvinfer1::IExecutionContext>(mCudaEngine->createExecutionContext(), InferDeleter());
assert(mContext != nullptr);
mContext->setOptimizationProfile(0);
std::cout << "Finished init" << std::endl;
}
void PSENet::detect(std::string image_path)
{
int batch_size = 1;
// Run inference
cv::Mat image = cv::imread(image_path);
int resize_h, resize_w;
float ratio_h, ratio_w;
auto start = std::chrono::system_clock::now();
float *input = preProcess(image, resize_h, resize_w, ratio_h, ratio_w);
float *output = new float[resize_h * resize_w * 6 / 16];
inferenceOnce(*mContext, input, output, resize_h, resize_w);
cv::Mat mask;
postProcess(output, mask, resize_h, resize_w);
drawRects(image, mask, ratio_h, ratio_w, stride_, 1.4);
auto end = std::chrono::system_clock::now();
cv::imwrite("result_" + image_path, image);
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
float *PSENet::preProcess(cv::Mat image, int &resize_h, int &resize_w, float &ratio_h, float &ratio_w)
{
cv::Mat imageRGB;
cv::cvtColor(image, imageRGB, CV_BGR2RGB);
cv::Mat imageProcessed;
int h = imageRGB.size().height;
int w = imageRGB.size().width;
resize_w = w;
resize_h = h;
float ratio = 1.0;
// limit the max side
if (resize_h > max_side_len_ && resize_w > max_side_len_)
{
if (resize_h > resize_w)
{
ratio = float(max_side_len_) / float(resize_h);
}
else
{
ratio = float(max_side_len_) / float(resize_w);
}
}
resize_h = int(resize_h * ratio);
resize_w = int(resize_w * ratio);
if (resize_h % 32 != 0)
{
resize_h = (resize_h / 32 + 1) * 32;
}
if (resize_w % 32 != 0)
{
resize_w = (resize_w / 32 + 1) * 32;
}
ratio_h = resize_h / float(h);
ratio_w = resize_w / float(w);
cv::resize(imageRGB, imageProcessed, cv::Size(resize_w, resize_h));
float *input = new float[3 * resize_h * resize_w];
cv::Mat imgFloat;
imageProcessed.convertTo(imgFloat, CV_32FC3);
cv::subtract(imgFloat, cv::Scalar(123.68, 116.78, 103.94), imgFloat, cv::noArray(), -1);
std::vector<cv::Mat> chw;
for (auto i = 0; i < 3; ++i)
{
chw.emplace_back(cv::Mat(cv::Size(resize_w, resize_h), CV_32FC1, input + i * resize_w * resize_h));
}
cv::split(imgFloat, chw);
return input;
}
void PSENet::postProcess(float *origin_output, cv::Mat &label_image, int resize_h, int resize_w)
{
// BxCxHxW S0 ===> S5 small ===> large
const int height = (resize_h + stride_ - 1) / stride_;
const int width = (resize_w + stride_ - 1) / stride_;
const int length = height * width;
std::vector<cv::Mat> kernels(num_kernels_);
cv::Mat max_kernel(height, width, CV_32F, (void *)(origin_output + (num_kernels_ - 1) * length), 0);
cv::threshold(max_kernel, max_kernel, post_threshold_, 255, cv::THRESH_BINARY);
max_kernel.convertTo(max_kernel, CV_8U);
assert(max_kernel.rows == height && max_kernel.cols == width);
for (auto i = 0; i < num_kernels_ - 1; ++i)
{
cv::Mat kernel = cv::Mat(height, width, CV_32F, (void *)(origin_output + i * length), 0);
cv::threshold(kernel, kernel, post_threshold_, 255, cv::THRESH_BINARY);
kernel.convertTo(kernel, CV_8U);
cv::bitwise_and(kernel, max_kernel, kernel);
assert(kernel.rows == height && kernel.cols == width);
kernels[i] = kernel;
}
kernels[num_kernels_ - 1] = max_kernel;
cv::Mat stats, centroids;
int num_labels = cv::connectedComponentsWithStats(kernels[0], label_image, stats, centroids, 4);
label_image.convertTo(label_image, CV_8U);
assert(label_image.rows == max_kernel.rows && label_image.cols == max_kernel.cols);
std::map<int, std::vector<cv::Point>> contourMaps;
// PSE algorithm
std::queue<std::tuple<int, int, int>> q;
std::queue<std::tuple<int, int, int>> q_next;
for (auto h = 0; h < height; ++h)
{
for (auto w = 0; w < width; ++w)
{
auto label = *label_image.ptr(h, w);
if (label > 0)
{
q.emplace(std::make_tuple(w, h, label));
contourMaps[label].emplace_back(cv::Point(w, h));
}
}
}
int dx[4] = {-1, 1, 0, 0};
int dy[4] = {0, 0, -1, 1};
for (auto idx = 1; idx < num_kernels_; ++idx)
{
auto *ptr_kernel = kernels[idx].data;
while (!q.empty())
{
auto q_n = q.front();
q.pop();
int x = std::get<0>(q_n);
int y = std::get<1>(q_n);
int l = std::get<2>(q_n);
bool is_edge = true;
for (auto j = 0; j < 4; ++j)
{
int tmpx = x + dx[j];
int tmpy = y + dy[j];
int offset = tmpy * width + tmpx;
if (tmpx < 0 || tmpx >= width || tmpy < 0 || tmpy >= height)
{
continue;
}
if (!(int)ptr_kernel[offset] || (int)*label_image.ptr(tmpy, tmpx) > 0)
{
continue;
}
q.emplace(std::make_tuple(tmpx, tmpy, l));
*label_image.ptr(tmpy, tmpx) = l;
contourMaps[l].emplace_back(cv::Point(tmpx, tmpy));
is_edge = false;
}
if (is_edge)
{
q_next.emplace(std::make_tuple(x, y, l));
}
}
std::swap(q, q_next);
}
}
#include "psenet.h"
#include <string>
#include <queue>
#define MAX_INPUT_SIZE 1200
#define MIN_INPUT_SIZE 128
#define OPT_INPUT_W 640
#define OPT_INPUT_H 640
PSENet::PSENet(int max_side_len, int min_side_len, float threshold, int num_kernel, int stride) : max_side_len_(max_side_len), min_side_len_(min_side_len),
post_threshold_(threshold),
num_kernels_(num_kernel),
stride_(stride)
{
}
PSENet::~PSENet()
{
}
// create the engine using only the API and not any parser.
ICudaEngine* PSENet::createEngine(IBuilder* builder, IBuilderConfig* config)
{
std::map<std::string, Weights> weightMap = loadWeights("./psenet.wts");
Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
ITensor* data = network->addInput(input_name_, dt, Dims4{ -1, 3, -1, -1 });
assert(data);
IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{ 7, 7 }, weightMap["resnet_v1_50/conv1/weights"], emptywts);
conv1->setStrideNd(DimsHW{ 2, 2 });
conv1->setPaddingNd(DimsHW{ 3, 3 });
assert(conv1);
IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "resnet_v1_50/conv1/BatchNorm/", 1e-5);
assert(bn1);
IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU);
assert(relu1);
// C2
IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 });
pool1->setStrideNd(DimsHW{ 2, 2 });
pool1->setPrePadding(DimsHW{ 0, 0 });
pool1->setPostPadding(DimsHW{ 1, 1 });
assert(pool1);
IActivationLayer* x;
x = bottleneck(network, weightMap, *pool1->getOutput(0), 64, 1, "resnet_v1_50/block1/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 64, 1, "resnet_v1_50/block1/unit_2/bottleneck_v1/", 0);
// C3
IActivationLayer* block1 = bottleneck(network, weightMap, *x->getOutput(0), 64, 2, "resnet_v1_50/block1/unit_3/bottleneck_v1/", 2);
x = bottleneck(network, weightMap, *block1->getOutput(0), 128, 1, "resnet_v1_50/block2/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 128, 1, "resnet_v1_50/block2/unit_2/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 128, 1, "resnet_v1_50/block2/unit_3/bottleneck_v1/", 0);
// C4
IActivationLayer* block2 = bottleneck(network, weightMap, *x->getOutput(0), 128, 2, "resnet_v1_50/block2/unit_4/bottleneck_v1/", 2);
x = bottleneck(network, weightMap, *block2->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_2/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_3/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_4/bottleneck_v1/", 0);
x = bottleneck(network, weightMap, *x->getOutput(0), 256, 1, "resnet_v1_50/block3/unit_5/bottleneck_v1/", 0);
IActivationLayer* block3 = bottleneck(network, weightMap, *x->getOutput(0), 256, 2, "resnet_v1_50/block3/unit_6/bottleneck_v1/", 2);
x = bottleneck(network, weightMap, *block3->getOutput(0), 512, 1, "resnet_v1_50/block4/unit_1/bottleneck_v1/", 1);
x = bottleneck(network, weightMap, *x->getOutput(0), 512, 1, "resnet_v1_50/block4/unit_2/bottleneck_v1/", 0);
// C5
IActivationLayer* block4 = bottleneck(network, weightMap, *x->getOutput(0), 512, 1, "resnet_v1_50/block4/unit_3/bottleneck_v1/", 0);
IActivationLayer* build_p5_r1 = addConvRelu(network, weightMap, *block4->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P5/");
assert(build_p5_r1);
IActivationLayer* build_p4_r1 = addConvRelu(network, weightMap, *block2->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P4/reduce_dimension/");
assert(build_p4_r1);
IResizeLayer* bfp_layer4_resize = network->addResize(*build_p5_r1->getOutput(0));
auto build_p4_r1_shape = network->addShape(*build_p4_r1->getOutput(0))->getOutput(0);
bfp_layer4_resize->setInput(1, *build_p4_r1_shape);
bfp_layer4_resize->setResizeMode(ResizeMode::kNEAREST);
bfp_layer4_resize->setAlignCorners(false);
assert(bfp_layer4_resize);
IElementWiseLayer* bfp_add = network->addElementWise(*bfp_layer4_resize->getOutput(0), *build_p4_r1->getOutput(0), ElementWiseOperation::kSUM);
assert(bfp_add);
IActivationLayer* build_p4_r2 = addConvRelu(network, weightMap, *bfp_add->getOutput(0), 256, 3, 1, "build_feature_pyramid/build_P4/avoid_aliasing/");
assert(build_p4_r2);
IActivationLayer* build_p3_r1 = addConvRelu(network, weightMap, *block1->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P3/reduce_dimension/");
assert(build_p3_r1);
IResizeLayer* bfp_layer3_resize = network->addResize(*build_p4_r2->getOutput(0));
bfp_layer3_resize->setResizeMode(ResizeMode::kNEAREST);
auto build_p3_r1_shape = network->addShape(*build_p3_r1->getOutput(0))->getOutput(0);
bfp_layer3_resize->setInput(1, *build_p3_r1_shape);
bfp_layer3_resize->setAlignCorners(false);
assert(bfp_layer3_resize);
IElementWiseLayer* bfp_add1 = network->addElementWise(*bfp_layer3_resize->getOutput(0), *build_p3_r1->getOutput(0), ElementWiseOperation::kSUM);
assert(bfp_add1);
IActivationLayer* build_p3_r2 = addConvRelu(network, weightMap, *bfp_add1->getOutput(0), 256, 3, 1, "build_feature_pyramid/build_P3/avoid_aliasing/");
assert(build_p3_r2);
IActivationLayer* build_p2_r1 = addConvRelu(network, weightMap, *pool1->getOutput(0), 256, 1, 1, "build_feature_pyramid/build_P2/reduce_dimension/");
assert(build_p2_r1);
IResizeLayer* bfp_layer2_resize = network->addResize(*build_p3_r2->getOutput(0));
bfp_layer2_resize->setResizeMode(ResizeMode::kNEAREST);
auto build_p2_r1_shape = network->addShape(*build_p2_r1->getOutput(0))->getOutput(0);
bfp_layer2_resize->setInput(1, *build_p2_r1_shape);
bfp_layer2_resize->setAlignCorners(false);
assert(bfp_layer2_resize);
IElementWiseLayer* bfp_add2 = network->addElementWise(*bfp_layer2_resize->getOutput(0), *build_p2_r1->getOutput(0), ElementWiseOperation::kSUM);
assert(bfp_add2);
// P2
IActivationLayer* build_p2_r2 = addConvRelu(network, weightMap, *bfp_add2->getOutput(0), 256, 3, 1, "build_feature_pyramid/build_P2/avoid_aliasing/");
assert(build_p2_r2);
auto build_p2_r2_shape = network->addShape(*build_p2_r2->getOutput(0))->getOutput(0);
// P3 x2
IResizeLayer* layer1_resize = network->addResize(*build_p3_r2->getOutput(0));
layer1_resize->setResizeMode(ResizeMode::kLINEAR);
layer1_resize->setInput(1, *build_p2_r2_shape);
layer1_resize->setAlignCorners(false);
assert(layer1_resize);
// P4 x4
IResizeLayer* layer2_resize = network->addResize(*build_p4_r2->getOutput(0));
layer2_resize->setResizeMode(ResizeMode::kLINEAR);
layer2_resize->setInput(1, *build_p2_r2_shape);
layer2_resize->setAlignCorners(false);
assert(layer2_resize);
// P5 x8
IResizeLayer* layer3_resize = network->addResize(*build_p5_r1->getOutput(0));
layer3_resize->setResizeMode(ResizeMode::kLINEAR);
layer3_resize->setInput(1, *build_p2_r2_shape);
layer3_resize->setAlignCorners(false);
assert(layer3_resize);
// C(P5,P4,P3,P2)
ITensor* inputTensors[] = { layer3_resize->getOutput(0), layer2_resize->getOutput(0), layer1_resize->getOutput(0), build_p2_r2->getOutput(0) };
IConcatenationLayer* concat = network->addConcatenation(inputTensors, 4);
assert(concat);
IConvolutionLayer* feature_result_conv = network->addConvolutionNd(*concat->getOutput(0), 256, DimsHW{ 3, 3 }, weightMap["feature_results/Conv/weights"], emptywts);
feature_result_conv->setPaddingNd(DimsHW{ 1, 1 });
assert(feature_result_conv);
IScaleLayer* feature_result_bn = addBatchNorm2d(network, weightMap, *feature_result_conv->getOutput(0), "feature_results/Conv/BatchNorm/", 1e-5);
assert(feature_result_bn);
IActivationLayer* feature_result_relu = network->addActivation(*feature_result_bn->getOutput(0), ActivationType::kRELU);
assert(feature_result_relu);
IConvolutionLayer* feature_result_conv_1 = network->addConvolutionNd(*feature_result_relu->getOutput(0), 6, DimsHW{ 1, 1 }, weightMap["feature_results/Conv_1/weights"], weightMap["feature_results/Conv_1/biases"]);
assert(feature_result_conv_1);
IActivationLayer* sigmoid = network->addActivation(*feature_result_conv_1->getOutput(0), ActivationType::kSIGMOID);
assert(sigmoid);
sigmoid->getOutput(0)->setName(output_name_);
std::cout << "Set name out" << std::endl;
network->markOutput(*sigmoid->getOutput(0));
// Set profile
IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions(input_name_, OptProfileSelector::kMIN, Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE));
profile->setDimensions(input_name_, OptProfileSelector::kOPT, Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W));
profile->setDimensions(input_name_, OptProfileSelector::kMAX, Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE));
config->addOptimizationProfile(profile);
// Build engine
config->setMaxWorkspaceSize(1 << 30); // 1G
#ifdef USE_FP16
config->setFlag(BuilderFlag::kFP16);
#endif
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
;
std::cout << "Build out" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*)(mem.second.values));
}
return engine;
}
void PSENet::serializeEngine()
{
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(builder, config);
assert(engine != nullptr);
// Serialize the engine
IHostMemory* modelStream{ nullptr };
modelStream = engine->serialize();
assert(modelStream != nullptr);
std::ofstream p("./psenet.engine", std::ios::binary | std::ios::out);
if (!p)
{
std::cerr << "Could not open plan output file" << std::endl;
return;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
return;
}
void PSENet::deserializeEngine()
{
std::ifstream file("./psenet.engine", std::ios::binary | std::ios::in);
if (file.good())
{
file.seekg(0, file.end);
size_t size = file.tellg();
file.seekg(0, file.beg);
char* trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
mCudaEngine = std::shared_ptr<nvinfer1::ICudaEngine>(mRuntime->deserializeCudaEngine(trtModelStream, size), InferDeleter());
assert(mCudaEngine != nullptr);
}
}
void PSENet::inferenceOnce(IExecutionContext& context, float* input, float* output, int input_h, int input_w)
{
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// 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(input_name_);
const int outputIndex = engine.getBindingIndex(output_name_);
context.setBindingDimensions(inputIndex, Dims4(1, 3, input_h, input_w));
int input_size = 3 * input_h * input_w * sizeof(float);
int output_size = input_h * input_w * 6 / 16 * sizeof(float);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], input_size));
CHECK(cudaMalloc(&buffers[outputIndex], output_size));
// 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, input_size, cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size, cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
void PSENet::init()
{
mRuntime = std::shared_ptr<nvinfer1::IRuntime>(createInferRuntime(gLogger), InferDeleter());
assert(mRuntime != nullptr);
std::cout << "Deserialize Engine" << std::endl;
deserializeEngine();
mContext = std::shared_ptr<nvinfer1::IExecutionContext>(mCudaEngine->createExecutionContext(), InferDeleter());
assert(mContext != nullptr);
mContext->setOptimizationProfile(0);
std::cout << "Finished init" << std::endl;
}
void PSENet::detect(std::string image_path)
{
// Run inference
cv::Mat image = cv::imread(image_path);
int resize_h, resize_w;
float ratio_h, ratio_w;
auto start = std::chrono::system_clock::now();
float* input = preProcess(image, resize_h, resize_w, ratio_h, ratio_w);
float* output = new float[resize_h * resize_w * 6 / 16];
inferenceOnce(*mContext, input, output, resize_h, resize_w);
std::vector<cv::RotatedRect> boxes = postProcess(output, resize_h, resize_w);
drawRects(image, boxes, stride_, ratio_h, ratio_w, 1.0);
auto end = std::chrono::system_clock::now();
cv::imwrite("result_" + image_path, image);
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
delete input;
delete output;
}
float* PSENet::preProcess(cv::Mat image, int& resize_h, int& resize_w, float& ratio_h, float& ratio_w)
{
cv::Mat imageRGB;
cv::cvtColor(image, imageRGB, cv::COLOR_BGR2RGB);
cv::Mat imageProcessed;
int h = imageRGB.size().height;
int w = imageRGB.size().width;
resize_w = w;
resize_h = h;
float ratio = 1.0;
// limit the max side and min side
if (resize_h > max_side_len_ || resize_w > max_side_len_)
{
if (resize_h > resize_w)
ratio = float(max_side_len_) / float(resize_h);
else
ratio = float(max_side_len_) / float(resize_w);
}
if (resize_h < min_side_len_ || resize_w < min_side_len_)
{
if (resize_h < resize_w)
ratio = float(min_side_len_) / float(resize_h);
else
ratio = float(min_side_len_) / float(resize_w);
}
resize_h = int(resize_h * ratio);
resize_w = int(resize_w * ratio);
if (resize_h % 32 != 0)
resize_h = (resize_h / 32 + 1) * 32;
if (resize_w % 32 != 0)
resize_w = (resize_w / 32 + 1) * 32;
ratio_h = resize_h / float(h);
ratio_w = resize_w / float(w);
cv::resize(imageRGB, imageProcessed, cv::Size(resize_w, resize_h));
float* input = new float[3 * resize_h * resize_w];
cv::Mat imgFloat;
imageProcessed.convertTo(imgFloat, CV_32FC3);
cv::subtract(imgFloat, cv::Scalar(123.68, 116.78, 103.94), imgFloat, cv::noArray(), -1);
std::vector<cv::Mat> chw;
for (auto i = 0; i < 3; ++i)
chw.emplace_back(cv::Mat(cv::Size(resize_w, resize_h), CV_32FC1, input + i * resize_w * resize_h));
cv::split(imgFloat, chw);
return input;
}
std::vector<cv::RotatedRect> PSENet::postProcess(float* origin_output, int resize_h, int resize_w)
{
// BxCxHxW S0 ===> S5 small ===> large
const int h = resize_h / stride_;
const int w = resize_w / stride_;
const int length = h * w;
// get kernels, sequence: 0->n, max -> min
std::vector<cv::Mat> kernels(num_kernels_);
for (auto i = num_kernels_ - 1; i >= 0; --i)
{
cv::Mat tmp_kernel(h, w, CV_32FC1, (void*)(origin_output + i * length), 0);
cv::threshold(tmp_kernel, tmp_kernel, post_threshold_, 255, cv::THRESH_BINARY);
tmp_kernel.convertTo(tmp_kernel, CV_8UC1);
assert(tmp_kernel.rows == h && tmp_kernel.cols == w);
kernels[num_kernels_ - 1 - i] = tmp_kernel;
}
cv::Mat stats, centroids, label_image;
int label_num = cv::connectedComponents(kernels[num_kernels_ - 1], label_image, 4);
label_image.convertTo(label_image, CV_8U);
assert(label_image.rows == h && label_image.cols == w);
cv::Mat out = cv::Mat::zeros(h, w, CV_8UC1);
std::queue<std::tuple<int, int, int>> q;
std::queue<std::tuple<int, int, int>> next_q;
for (int i = 0; i < h; i++)
{
for (int j = 0; j < w; j++)
{
auto label = *label_image.ptr(i, j);
if (label > 0)
{
q.push(std::make_tuple(i, j, label));
*out.ptr(i, j) = label;
}
}
}
int dx[4] = { -1, 1, 0, 0 };
int dy[4] = { 0, 0, -1, 1 };
for (int i = num_kernels_ - 2; i >= 0; i--)
{
//get each kernels
auto kernel = kernels[i];
while (!q.empty())
{
//get each queue menber in q
auto q_n = q.front();
q.pop();
int y = std::get<0>(q_n); //i
int x = std::get<1>(q_n); //j
int l = std::get<2>(q_n); //label
//store the edge pixel after one expansion
bool is_edge = true;
for (int idx = 0; idx < 4; idx++)
{
int index_y = y + dy[idx];
int index_x = x + dx[idx];
if (index_y < 0 || index_y >= h || index_x < 0 || index_x >= w)
continue;
if (!*kernel.ptr(index_y, index_x) || *out.ptr(index_y, index_x) > 0)
continue;
q.push(std::make_tuple(index_y, index_x, l));
*out.ptr(index_y, index_x) = l;
is_edge = false;
}
if (is_edge)
{
next_q.push(std::make_tuple(y, x, l));
}
}
std::swap(q, next_q);
}
std::vector<cv::RotatedRect> boxes;
for (auto n = 1; n < label_num; ++n)
{
std::vector<cv::Point> points;
cv::findNonZero(out == n, points);
cv::Mat fuck = out == n;
cv::RotatedRect rect = cv::minAreaRect(points);
boxes.emplace_back(rect);
}
return boxes;
}

78
psenet/psenet.h
View File

@ -1,39 +1,39 @@
#ifndef TENSORRTX_PSENET_H
#define TENSORRTX_PSENET_H
#include <memory>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include "utils.h"
#include "layers.h"
class PSENet
{
public:
PSENet(int max_side_len, float threshold, int num_kernel, int stride);
~PSENet();
ICudaEngine *createEngine(IBuilder *builder, IBuilderConfig *config);
void serializeEngine();
void deserializeEngine();
void init();
void inferenceOnce(IExecutionContext &context, float *input, float *output, int input_h, int input_w);
void detect(std::string image_path);
float *preProcess(cv::Mat image, int &resize_h, int &resize_w, float &ratio_h, float &ratio_w);
void postProcess(float *origin_output, cv::Mat &label_image, int resize_h, int resize_w);
private:
Logger gLogger;
std::shared_ptr<nvinfer1::IRuntime> mRuntime;
std::shared_ptr<nvinfer1::ICudaEngine> mCudaEngine;
std::shared_ptr<nvinfer1::IExecutionContext> mContext;
DataType dt = DataType::kFLOAT;
const char *input_name_ = "input";
const char *output_name_ = "maps";
int max_side_len_ = 640;
float post_threshold_ = 0.9;
int num_kernels_ = 6;
int stride_ = 4;
};
#endif // TENSORRTX_PSENET_H
#ifndef TENSORRTX_PSENET_H
#define TENSORRTX_PSENET_H
#include <memory>
#include <vector>
#include <chrono>
#include <opencv2/opencv.hpp>
#include "utils.h"
#include "layers.h"
class PSENet
{
public:
PSENet(int max_side_len, int min_side_len, float threshold, int num_kernel, int stride);
~PSENet();
ICudaEngine* createEngine(IBuilder* builder, IBuilderConfig* config);
void serializeEngine();
void deserializeEngine();
void init();
void inferenceOnce(IExecutionContext& context, float* input, float* output, int input_h, int input_w);
void detect(std::string image_path);
float* preProcess(cv::Mat image, int& resize_h, int& resize_w, float& ratio_h, float& ratio_w);
std::vector<cv::RotatedRect> postProcess(float* origin_output, int resize_h, int resize_w);
private:
Logger gLogger;
std::shared_ptr<nvinfer1::IRuntime> mRuntime;
std::shared_ptr<nvinfer1::ICudaEngine> mCudaEngine;
std::shared_ptr<nvinfer1::IExecutionContext> mContext;
DataType dt = DataType::kFLOAT;
const char* input_name_ = "input";
const char* output_name_ = "output";
int max_side_len_ = 1024;
int min_side_len_ = 640;
float post_threshold_ = 0.9;
int num_kernels_ = 6;
int stride_ = 4;
};
#endif // TENSORRTX_PSENET_H

141
psenet/utils.cpp
View File

@ -1,73 +1,68 @@
#include "utils.h"
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::cout << "Model weight is large, it will take some time." << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
std::cout << "Finish load weight" << std::endl;
return weightMap;
}
cv::RotatedRect expandBox(const cv::RotatedRect &inBox, float ratio)
{
cv::Size size = inBox.size;
int neww = int(size.width * ratio);
int newh = int(size.height * ratio);
return cv::RotatedRect(inBox.center, cv::Size(neww, newh), inBox.angle);
}
void drawRects(cv::Mat &image, cv::Mat mask, float ratio_h, float ratio_w, int stride, float expand_ratio)
{
std::vector<std::vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarcy;
cv::findContours(mask, 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 (auto i = 0; i < contours.size(); i++)
{
box[i] = cv::minAreaRect(cv::Mat(contours[i]));
cv::RotatedRect expandbox = expandBox(box[i], expand_ratio);
expandbox.points(rect);
for (auto j = 0; j < 4; j++)
{
cv::line(image, cv::Point{int(rect[j].x / ratio_w * stride), int(rect[j].y / ratio_h * stride)}, cv::Point{int(rect[(j + 1) % 4].x / ratio_w * stride), int(rect[(j + 1) % 4].y / ratio_h * stride)}, cv::Scalar(0, 0, 255), 2, 8);
}
}
}
#include "utils.h"
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::cout << "Model weight is large, it will take some time." << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{ DataType::kFLOAT, nullptr, 0 };
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
std::cout << "Finish load weight" << std::endl;
return weightMap;
}
cv::RotatedRect expandBox(const cv::RotatedRect& inBox, float ratio)
{
cv::Size size = inBox.size;
int neww = int(size.width * ratio);
int newh = int(size.height * ratio);
return cv::RotatedRect(inBox.center, cv::Size(neww, newh), inBox.angle);
}
void drawRects(cv::Mat& image, std::vector<cv::RotatedRect> boxes, float stride, float ratio_h, float ratio_w, float expand_ratio)
{
cv::Point2f rect[4];
for (unsigned int i = 0; i < boxes.size(); i++)
{
cv::RotatedRect box = boxes[i];
cv::RotatedRect expandbox = expandBox(box, expand_ratio);
expandbox.points(rect);
for (auto j = 0; j < 4; j++)
{
cv::line(image, cv::Point{ int(rect[j].x / ratio_w * stride), int(rect[j].y / ratio_h * stride) }, cv::Point{ int(rect[(j + 1) % 4].x / ratio_w * stride), int(rect[(j + 1) % 4].y / ratio_h * stride) }, cv::Scalar(0, 0, 255), 2, 8);
}
}
}

165
psenet/utils.h
View File

@ -1,82 +1,83 @@
#ifndef TENSORRTX_UTILS_H
#define TENSORRTX_UTILS_H
#include <map>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "assert.h"
using namespace nvinfer1;
std::map<std::string, Weights> loadWeights(const std::string file);
cv::RotatedRect expandBox(const cv::RotatedRect &inBox, float ratio = 1.0);
void drawRects(cv::Mat &image, cv::Mat mask, float ratio_h, float ratio_w, int stride, float expand_ratio = 1.4);
cv::Mat renderSegment(cv::Mat image, const cv::Mat &mask);
// <============== Operator =============>
struct InferDeleter
{
template <typename T>
void operator()(T *obj) const
{
if (obj)
{
obj->destroy();
}
}
};
#define CHECK(status) \
do \
{ \
auto ret = (status); \
if (ret != 0) \
{ \
std::cout << "Cuda failure: " << ret; \
abort(); \
} \
} while (0)
// Logger for TensorRT info/warning/errors
class Logger : public nvinfer1::ILogger
{
public:
Logger() : Logger(Severity::kWARNING) {}
Logger(Severity severity) : reportableSeverity(severity) {}
void log(Severity severity, const char *msg) override
{
// suppress messages with severity enum value greater than the reportable
if (severity > reportableSeverity)
return;
switch (severity)
{
case Severity::kINTERNAL_ERROR:
std::cerr << "INTERNAL_ERROR: ";
break;
case Severity::kERROR:
std::cerr << "ERROR: ";
break;
case Severity::kWARNING:
std::cerr << "WARNING: ";
break;
case Severity::kINFO:
std::cerr << "INFO: ";
break;
default:
std::cerr << "UNKNOWN: ";
break;
}
std::cerr << msg << std::endl;
}
Severity reportableSeverity{Severity::kWARNING};
};
#endif
#ifndef TENSORRTX_UTILS_H
#define TENSORRTX_UTILS_H
#include <map>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "assert.h"
#include <fstream>
using namespace nvinfer1;
std::map<std::string, Weights> loadWeights(const std::string file);
cv::RotatedRect expandBox(const cv::RotatedRect& inBox, float ratio = 1.0);
void drawRects(cv::Mat& image, std::vector<cv::RotatedRect> boxes, float stride, float ratio_h, float ratio_w, float expand_ratio);
cv::Mat renderSegment(cv::Mat image, const cv::Mat& mask);
// <============== Operator =============>
struct InferDeleter
{
template <typename T>
void operator()(T* obj) const
{
if (obj)
{
obj->destroy();
}
}
};
#define CHECK(status) \
do \
{ \
auto ret = (status); \
if (ret != 0) \
{ \
std::cout << "Cuda failure: " << ret; \
abort(); \
} \
} while (0)
// Logger for TensorRT info/warning/errors
class Logger : public nvinfer1::ILogger
{
public:
Logger() : Logger(Severity::kWARNING) {}
Logger(Severity severity) : reportableSeverity(severity) {}
void log(Severity severity, const char* msg) override
{
// suppress messages with severity enum value greater than the reportable
if (severity > reportableSeverity)
return;
switch (severity)
{
case Severity::kINTERNAL_ERROR:
std::cerr << "INTERNAL_ERROR: ";
break;
case Severity::kERROR:
std::cerr << "ERROR: ";
break;
case Severity::kWARNING:
std::cerr << "WARNING: ";
break;
case Severity::kINFO:
std::cerr << "INFO: ";
break;
default:
std::cerr << "UNKNOWN: ";
break;
}
std::cerr << msg << std::endl;
}
Severity reportableSeverity{ Severity::kWARNING };
};
#endif