diff --git a/psenet/README.md b/psenet/README.md
index 93d2efe..a6b9c6c 100644
--- a/psenet/README.md
+++ b/psenet/README.md
@@ -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.
+
-
-
## 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.
-
diff --git a/psenet/gen_tf_wts.py b/psenet/gen_tf_wts.py
index 0501b75..2d94c0c 100644
--- a/psenet/gen_tf_wts.py
+++ b/psenet/gen_tf_wts.py
@@ -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")
\ No newline at end of file
diff --git a/psenet/layers.cpp b/psenet/layers.cpp
index acbaba0..7937a34 100644
--- a/psenet/layers.cpp
+++ b/psenet/layers.cpp
@@ -1,136 +1,110 @@
-#include "layers.h"
-
-IScaleLayer *addBatchNorm2d(INetworkDefinition *network, std::map &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(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(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(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 &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 &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& 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(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(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(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& 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& 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;
}
\ No newline at end of file
diff --git a/psenet/main.cpp b/psenet/main.cpp
index af83730..1dfdbfa 100644
--- a/psenet/main.cpp
+++ b/psenet/main.cpp
@@ -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 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 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;
+ }
+}
diff --git a/psenet/psenet.cpp b/psenet/psenet.cpp
index 85be8e9..504cf6c 100644
--- a/psenet/psenet.cpp
+++ b/psenet/psenet.cpp
@@ -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 weightMap = loadWeights("./psenet.wts");
- Weights emptywts{DataType::kFLOAT, nullptr, 0};
- const auto explicitBatch = 1U << static_cast(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(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(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(createInferRuntime(gLogger), InferDeleter());
- assert(mRuntime != nullptr);
-
- std::cout << "Deserialize Engine" << std::endl;
- deserializeEngine();
-
- mContext = std::shared_ptr(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(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 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 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> contourMaps;
-
- // PSE algorithm
- std::queue> q;
- std::queue> 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
+#include
+#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 weightMap = loadWeights("./psenet.wts");
+ Weights emptywts{ DataType::kFLOAT, nullptr, 0 };
+ const auto explicitBatch = 1U << static_cast(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(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(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(createInferRuntime(gLogger), InferDeleter());
+ assert(mRuntime != nullptr);
+
+ std::cout << "Deserialize Engine" << std::endl;
+ deserializeEngine();
+
+ mContext = std::shared_ptr(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 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(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 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 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 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> q;
+ std::queue> 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 boxes;
+ for (auto n = 1; n < label_num; ++n)
+ {
+ std::vector points;
+ cv::findNonZero(out == n, points);
+ cv::Mat fuck = out == n;
+ cv::RotatedRect rect = cv::minAreaRect(points);
+ boxes.emplace_back(rect);
+ }
+ return boxes;
+}
diff --git a/psenet/psenet.h b/psenet/psenet.h
index fe44369..e0e290f 100644
--- a/psenet/psenet.h
+++ b/psenet/psenet.h
@@ -1,39 +1,39 @@
-#ifndef TENSORRTX_PSENET_H
-#define TENSORRTX_PSENET_H
-#include
-#include
-#include
-#include
-#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 mRuntime;
- std::shared_ptr mCudaEngine;
- std::shared_ptr 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
+#include
+#include
+#include
+#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 postProcess(float* origin_output, int resize_h, int resize_w);
+
+private:
+ Logger gLogger;
+ std::shared_ptr mRuntime;
+ std::shared_ptr mCudaEngine;
+ std::shared_ptr 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
diff --git a/psenet/utils.cpp b/psenet/utils.cpp
index c66cb33..5584646 100644
--- a/psenet/utils.cpp
+++ b/psenet/utils.cpp
@@ -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]
-std::map 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 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(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> contours;
- std::vector hierarcy;
- cv::findContours(mask, contours, hierarcy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);
-
- std::vector boundRect(contours.size());
- std::vector 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]
+std::map 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 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(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 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);
+ }
+ }
+}
diff --git a/psenet/utils.h b/psenet/utils.h
index 77d99d0..e6430c7 100644
--- a/psenet/utils.h
+++ b/psenet/utils.h
@@ -1,82 +1,83 @@
-#ifndef TENSORRTX_UTILS_H
-#define TENSORRTX_UTILS_H
-
-#include