From 9071fa2d2b9f5d3f714e614979656b3efd8024aa Mon Sep 17 00:00:00 2001 From: wang-xinyu Date: Fri, 25 Nov 2022 21:42:01 +0800 Subject: [PATCH] unet fix coding style --- unet/common.hpp | 8 -------- unet/unet.cpp | 54 ++++++++++++++++++++++++------------------------- 2 files changed, 26 insertions(+), 36 deletions(-) diff --git a/unet/common.hpp b/unet/common.hpp index 3121f58..ac2ec89 100644 --- a/unet/common.hpp +++ b/unet/common.hpp @@ -9,7 +9,6 @@ #include #include "NvInfer.h" - #define CHECK(status) \ do\ {\ @@ -23,10 +22,6 @@ using namespace nvinfer1; - - - - // TensorRT weight files have a simple space delimited format: // [type] [size] std::map loadWeights(const std::string file) { @@ -100,7 +95,6 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) { Weights emptywts{DataType::kFLOAT, nullptr, 0}; int p = ksize / 2; @@ -121,8 +115,6 @@ ILayer* convBlock(INetworkDefinition *network, std::map& w return ew; } - - int read_files_in_dir(const char *p_dir_name, std::vector &file_names) { DIR *p_dir = opendir(p_dir_name); if (p_dir == nullptr) { diff --git a/unet/unet.cpp b/unet/unet.cpp index 702bf03..a075f2e 100644 --- a/unet/unet.cpp +++ b/unet/unet.cpp @@ -64,7 +64,6 @@ ILayer* doubleConv(INetworkDefinition *network, std::map& } ILayer* down(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int p, std::string lname){ - IPoolingLayer* pool1 = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{2, 2}); assert(pool1); ILayer* dcov1 = doubleConv(network,weightMap,*pool1->getOutput(0),outch,3,lname+".maxpool_conv.1",outch); @@ -119,7 +118,7 @@ ILayer* outConv(INetworkDefinition *network, std::map& wei return conv1; } -ICudaEngine* createEngine_l(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { +ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { INetworkDefinition* network = builder->createNetworkV2(0U); // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME @@ -130,16 +129,16 @@ ICudaEngine* createEngine_l(unsigned int maxBatchSize, IBuilder* builder, IBuild Weights emptywts{DataType::kFLOAT, nullptr, 0}; // build network - auto x1 = doubleConv(network,weightMap,*data,64,3,"inc",64); - auto x2 = down(network,weightMap,*x1->getOutput(0),128,1,"down1"); - auto x3 = down(network,weightMap,*x2->getOutput(0),256,1,"down2"); - auto x4 = down(network,weightMap,*x3->getOutput(0),512,1,"down3"); - auto x5 = down(network,weightMap,*x4->getOutput(0),512,1,"down4"); - ILayer* x6 = up(network,weightMap,*x5->getOutput(0),*x4->getOutput(0),512,512,512,"up1"); - ILayer* x7 = up(network,weightMap,*x6->getOutput(0),*x3->getOutput(0),256,256,256,"up2"); - ILayer* x8 = up(network,weightMap,*x7->getOutput(0),*x2->getOutput(0),128,128,128,"up3"); - ILayer* x9 = up(network,weightMap,*x8->getOutput(0),*x1->getOutput(0),64,64,64,"up4"); - ILayer* x10 = outConv(network,weightMap,*x9->getOutput(0),OUTPUT_SIZE,"outc"); + auto x1 = doubleConv(network, weightMap, *data, 64, 3, "inc", 64); + auto x2 = down(network, weightMap, *x1->getOutput(0), 128, 1, "down1"); + auto x3 = down(network, weightMap, *x2->getOutput(0), 256, 1, "down2"); + auto x4 = down(network, weightMap, *x3->getOutput(0), 512, 1, "down3"); + auto x5 = down(network, weightMap, *x4->getOutput(0), 512, 1, "down4"); + ILayer* x6 = up(network, weightMap, *x5->getOutput(0), *x4->getOutput(0), 512, 512, 512, "up1"); + ILayer* x7 = up(network, weightMap, *x6->getOutput(0), *x3->getOutput(0), 256, 256, 256, "up2"); + ILayer* x8 = up(network, weightMap, *x7->getOutput(0), *x2->getOutput(0), 128, 128, 128, "up3"); + ILayer* x9 = up(network, weightMap, *x8->getOutput(0), *x1->getOutput(0), 64, 64, 64, "up4"); + ILayer* x10 = outConv(network, weightMap, *x9->getOutput(0), OUTPUT_SIZE, "outc"); std::cout << "set name out" << std::endl; x10->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*x10->getOutput(0)); @@ -171,7 +170,7 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { IBuilderConfig* config = builder->createBuilderConfig(); // Create model to populate the network, then set the outputs and create an engine - ICudaEngine* engine = createEngine_l(maxBatchSize, builder, config, DataType::kFLOAT); + ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT); assert(engine != nullptr); // Serialize the engine @@ -216,15 +215,15 @@ void doInference(IExecutionContext& context, float* input, float* output, int ba CHECK(cudaFree(buffers[outputIndex])); } -struct Detection { - float mask[INPUT_W*INPUT_H*1]; +struct Detection { + float mask[INPUT_W * INPUT_H * 1]; }; float sigmoid(float x) { return (1 / (1 + exp(-x))); } -void process_cls_result(Detection &res, float *output) { +void process_cls_result(Detection &res, float *output) { for (int i = 0; i < INPUT_W * INPUT_H * 1; i++) { res.mask[i] = sigmoid(*(output+i)); } @@ -293,7 +292,7 @@ int main(int argc, char** argv) { cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]); if (img.empty()) continue; cv::Mat pr_img = preprocess_img(img); // letterbox BGR to RGB - // cv::imwrite("s_o" + file_names[f - fcount + 1 + b] + "_unet.jpg", pr_img); + // cv::imwrite("s_o" + file_names[f - fcount + 1 + b] + "_unet.jpg", pr_img); int i = 0; for (int row = 0; row < INPUT_H; ++row) { uchar* uc_pixel = pr_img.data + row * pr_img.step; @@ -323,24 +322,23 @@ int main(int argc, char** argv) { for (int b = 0; b < fcount; b++) { auto& res = batch_res[b]; - float * mask = res.mask; - cv::Mat mask_mat = cv::Mat(INPUT_H,INPUT_W,CV_8UC1); - uchar *ptmp = NULL; - for(int i =0; i< INPUT_H ;i++){ + float* mask = res.mask; + cv::Mat mask_mat = cv::Mat(INPUT_H, INPUT_W, CV_8UC1); + uchar* ptmp = NULL; + for (int i = 0; i < INPUT_H; i++) { ptmp = mask_mat.ptr(i); - for(int j=0;j CONF_THRESH){ + if (*pixcel > CONF_THRESH) { ptmp[j] = 255; - } - else{ + } else { ptmp[j]=0; } - } + } } - cv::imwrite("s_" + file_names[f - fcount + 1 + b] + "_unet.jpg", mask_mat); + cv::imwrite("s_" + file_names[f - fcount + 1 + b] + "_unet.jpg", mask_mat); } fcount = 0; }