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

unet fix coding style

Browse Source
This commit is contained in:
wang-xinyu 2022-11-25 21:42:01 +08:00
parent e061888ed2
commit 9071fa2d2b
2 changed files with 26 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
unet/common.hpp
View File

@ -9,7 +9,6 @@
#include <dirent.h>
#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] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file) {
@ -100,7 +95,6 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, W
return scale_1;
}
ILayer* convBlock(INetworkDefinition *network, std::map<std::string, Weights>& 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<std::string, Weights>& w
return ew;
}
int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
DIR *p_dir = opendir(p_dir_name);
if (p_dir == nullptr) {

54
unet/unet.cpp
View File

@ -64,7 +64,6 @@ ILayer* doubleConv(INetworkDefinition *network, std::map<std::string, Weights>&
}
ILayer* down(INetworkDefinition *network, std::map<std::string, Weights>& 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<std::string, Weights>& 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<uchar>(i);
for(int j=0;j<INPUT_W;j++){
float * pixcel = mask+i*INPUT_W+j;
for (int j = 0; j < INPUT_W; j++){
float * pixcel = mask + i * INPUT_W + j;
// std::cout << *pixcel << std::endl;
if(*pixcel > 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;
}