From 79cd96de6fea52d3e6c9fb0a710dde2910c16624 Mon Sep 17 00:00:00 2001 From: Micaela Verucchi Date: Thu, 2 Jul 2020 12:14:20 +0200 Subject: [PATCH] Add resize to original size, writing of segmentation Signed-off-by: Micaela Verucchi --- demo/demo/seg_demo.cpp | 39 ++++++++++++++++--- include/tkDNN/SegmentationNN.h | 71 +++++++++++++++++++++------------- 2 files changed, 79 insertions(+), 31 deletions(-) diff --git a/demo/demo/seg_demo.cpp b/demo/demo/seg_demo.cpp index 0737ad7..22ac2d1 100644 --- a/demo/demo/seg_demo.cpp +++ b/demo/demo/seg_demo.cpp @@ -14,6 +14,24 @@ void sig_handler(int signo) { gRun = false; } +void writePred(const std::string& images_names, const std::string& gt_folder, const std::string& out_folder, tk::dnn::SegmentationNN& segNN){ + std::ifstream all_gt(images_names); + std::string filename; + cv::Mat frame; + std::vector batch_frame; + std::vector batch_dnn_input; + for (; std::getline(all_gt, filename); ) { + std::cout< 5) show = atoi(argv[5]); + bool write_pred = false; + if(argc > 6) + write_pred = atoi(argv[6]); + if(n_batch < 1 || n_batch > 64) FatalError("Batch dim not supported"); - if(!show) - SAVE_RESULT = true; - - - tk::dnn::SegmentationNN segNN; segNN.init(net, n_classes, n_batch); + if(write_pred){ + std::string gt_folder = "../demo/CityScapes_val/images/"; + std::string images_names = "../demo/CityScapes_val/all_images.txt"; + std::string out_folder = "seg/"; + + writePred(images_names, gt_folder, out_folder, segNN); + return 0; + } + + if(!show) + SAVE_RESULT = true; + gRun = true; cv::VideoCapture cap(input); diff --git a/include/tkDNN/SegmentationNN.h b/include/tkDNN/SegmentationNN.h index 87c7b8c..927feb0 100644 --- a/include/tkDNN/SegmentationNN.h +++ b/include/tkDNN/SegmentationNN.h @@ -26,7 +26,6 @@ class SegmentationNN { int nBatches = 1; std::vector originalSize; - std::vector masks; cv::Mat bgr[3]; dnnType *input; dnnType *input_d; @@ -40,6 +39,24 @@ class SegmentationNN { cublasHandle_t cublasHandle; + void computeBorders(const int or_width, const int or_height, int& top, int& bottom, int& left, int&right){ + top = 0; + bottom = 0; + left = 0; + right = 0; + + if(or_height != or_width){ + if(or_height < or_width){ + top = (or_width - or_height)/2; + bottom = or_width - top - or_height; + } + else{ + left = (or_height - or_width)/2; + right = or_height - left - or_width; + } + } + } + /** * This method preprocess the image, before feeding it to the NN. * @@ -47,32 +64,20 @@ class SegmentationNN { * @param bi batch index */ void preprocess(cv::Mat &frame, const int bi=0) { + originalSize[bi] = frame.size(); + frame.convertTo(frame, CV_32FC3, 1 / 255.0, 0); int H = frame.rows; int W = frame.cols; cv::Mat frame_cropped; - cv::Mat mask(frame.size(), CV_8UC3, cv::Scalar(255,255,255)); - - if(H != W){ - if(H < W){ - int top = (W - H)/2; - int bottom = W - top - H; - cv::copyMakeBorder(frame, frame_cropped, top, bottom, 0, 0, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) ); - cv::copyMakeBorder(mask, mask, top, bottom, 0, 0, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) ); - } - else{ - int left = (H - W)/2; - int right = H - left - W; - cv::copyMakeBorder(frame, frame_cropped, 0, 0, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) ); - cv::copyMakeBorder(mask, mask, 0, 0, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) ); - } - } + + int top, bottom, left, right; + computeBorders(W, H, top, bottom, left, right); + cv::copyMakeBorder(frame, frame_cropped, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) ); tk::dnn::dataDim_t idim = netRT->input_dim; resize(frame_cropped, frame_cropped, cv::Size(idim.w, idim.h)); - resize(mask, mask, cv::Size(idim.w, idim.h)); - masks[bi] = mask; cv::split(frame_cropped, bgr); for (int i = 0; i < idim.c; i++){ @@ -92,7 +97,7 @@ class SegmentationNN { * * @param bi batch index */ - void postprocess(const int bi=0) { + void postprocess(const int bi=0, bool appy_colormap = true) { dnnType *rt_out = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi; dataDim_t odim = netRT->output_dim; @@ -103,7 +108,23 @@ class SegmentationNN { dataDim_t vdim = odim; vdim.c = 1; - segmented[bi] = vizData2Mat(tmpOutData_h, vdim, 1024, 0, 18); + + cv::Mat colored; + + if(appy_colormap) + colored = vizData2Mat(tmpOutData_h, vdim, 1024, 0, 18); + else{ + cv::Mat colored_fp32 (cv::Size(odim.w, odim.h),CV_32FC1, tmpOutData_h); + colored_fp32.convertTo(colored, CV_8UC1); + } + + int max_dim = (originalSize[bi].width > originalSize[bi].height) ? originalSize[bi].width : originalSize[bi].height; + resize(colored, colored, cv::Size(max_dim, max_dim)); + int top, bottom, left, right; + computeBorders(originalSize[bi].width, originalSize[bi].height, top, bottom, left, right); + cv::Rect roi(left,top,originalSize[bi].width, originalSize[bi].height); + cv::Mat or_size (colored, roi); + segmented[bi] = or_size; }; public: @@ -148,7 +169,7 @@ class SegmentationNN { checkCuda(cudaMallocHost(&tmpOutData_h, sizeof(float) * odim.w*odim.h)); segmented.resize(nBatches); - masks.resize(nBatches); + originalSize.resize(nBatches); std::vector mean = {0.485, 0.456, 0.406}; std::vector stddev = {0.229, 0.224, 0.225}; @@ -172,7 +193,7 @@ class SegmentationNN { * @param mAP set to true only if all the probabilities for a bounding * box are needed, as in some cases for the mAP calculation */ - void update(std::vector& frames, const int cur_batches=1){ + void update(std::vector& frames, const int cur_batches=1, bool apply_colormap=true){ if(cur_batches > nBatches) FatalError("A batch size greater than nBatches cannot be used"); @@ -204,7 +225,7 @@ class SegmentationNN { { TKDNN_TSTART for(int bi=0; bi