Post-processing improved

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>

demo/demo/seg_demo.cpp

@@ -7,7 +7,7 @@
 #include "SegmentationNN.h"
 
 bool gRun;
-bool SAVE_RESULT = false;
+bool SAVE_RESULT = true;
 
 void sig_handler(int signo) {
     std::cout<<"request gateway stop\n";
@@ -56,13 +56,10 @@ int main(int argc, char *argv[]) {
     if(SAVE_RESULT) {
         int w = cap.get(cv::CAP_PROP_FRAME_WIDTH);
         int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
-        resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M','P','4','V'), 30, cv::Size(w, h));
+        resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M','P','4','V'), 30, cv::Size(1024, 1024));
     }
 
     cv::Mat frame;
-    if(show)
-        cv::namedWindow("segmentation", cv::WINDOW_NORMAL);
-
     std::vector<cv::Mat> batch_frame;
     std::vector<cv::Mat> batch_dnn_input;
 
@@ -85,14 +82,8 @@ int main(int argc, char *argv[]) {
     
         //inference
         segNN.update(batch_dnn_input, n_batch);
-        segNN.draw();
+        frame = segNN.draw();
 
-        if(show){
-            for(int bi=0; bi< n_batch; ++bi){
-                cv::imshow("segmentation", batch_frame[bi]);
-                cv::waitKey(1);
-            }
-        }
         if(n_batch == 1 && SAVE_RESULT)
             resultVideo << frame;
     }

include/tkDNN/SegmentationNN.h

@@ -17,6 +17,7 @@
 
 #include "tkdnn.h"
 #include "NetworkViz.h"
+#include "kernelsThrust.h"
 
 namespace tk { namespace dnn {
 
@@ -33,6 +34,12 @@ class SegmentationNN {
         dnnType *input_d;
         float* confidences_h;
 
+        float * tmpInputData_d;
+        float *tmpOutData_d;
+        float *tmpOutData_h;
+
+        cublasHandle_t cublasHandle;
+
         /**
          * This method preprocess the image, before feeding it to the NN.
          *
@@ -97,28 +104,14 @@ class SegmentationNN {
             dnnType *rt_out = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;
 
             dataDim_t odim = netRT->output_dim;
-
             
-            checkCuda(cudaMemcpy(confidences_h, rt_out, odim.tot() * sizeof(float), cudaMemcpyDeviceToHost));
+            matrixTranspose(cublasHandle, rt_out, tmpInputData_d, odim.c, odim.w*odim.h);
+            maxElem(tmpInputData_d, tmpOutData_d, odim.c, odim.h, odim.w);
+            checkCuda(cudaMemcpy(tmpOutData_h, tmpOutData_d, odim.w*odim.h * sizeof(float), cudaMemcpyDeviceToHost));
 
-            for(int i=0;i<odim.h;++i){
-                for(int j=0;j<odim.w;++j){
-                    float max_conf = 0;
-                    int max_id = 0;
-
-                    for(int k=0; k<odim.c;++k){
-                        float cur_conf = confidences_h[bi*odim.tot()+k*odim.h*odim.w+i*odim.h+j];
-                        if(cur_conf > max_conf){
-                            max_conf = cur_conf;
-                            max_id = k;
-                        }
-                    }
-                    confidences_h[bi*odim.tot()+0*odim.h*odim.w+i*odim.h+j] = max_id;
-                }
-            }
             dataDim_t vdim = odim;
             vdim.c = 1;
-            segmented[bi] = vizData2Mat(confidences_h, vdim, 1024, 0, 18);
+            segmented[bi] = vizData2Mat(tmpOutData_h, vdim, 1024, 0, 18);
         };
 
     public:
@@ -127,8 +120,12 @@ class SegmentationNN {
         std::vector<std::string> classesNames;
         std::vector<cv::Mat> segmented;
 
-        SegmentationNN() {};
+        SegmentationNN() {
-        ~SegmentationNN(){};
+            checkERROR( cublasCreate(&cublasHandle) );
+        };
+        ~SegmentationNN(){
+            checkERROR( cublasDestroy(cublasHandle) );
+        };
 
         /**
          * Method used to inialize the class, allocate memory and compute 
@@ -151,7 +148,12 @@ class SegmentationNN {
             checkCuda(cudaMallocHost(&input, sizeof(dnnType) * netRT->input_dim.tot() * nBatches));
             checkCuda(cudaMalloc(&input_d, sizeof(dnnType) * netRT->input_dim.tot() * nBatches));
 
-            confidences_h = (float *)malloc(netRT->output_dim.tot() * sizeof(float));
+            dataDim_t odim = netRT->output_dim;
+
+            checkCuda(cudaMallocHost(&confidences_h, sizeof(float) * odim.tot()));
+            checkCuda(cudaMalloc(&tmpInputData_d, sizeof(float) * odim.tot()));
+            checkCuda(cudaMalloc(&tmpOutData_d, sizeof(float) * odim.w*odim.h));
+            checkCuda(cudaMallocHost(&tmpOutData_h, sizeof(float) * odim.w*odim.h));
 
             segmented.resize(nBatches);
             masks.resize(nBatches);
@@ -208,7 +210,7 @@ class SegmentationNN {
         /**
          * Method to draw boundixg boxes and labels on a frame.
          */
-        void draw(const int cur_batches=1) {
+        cv::Mat draw(const int cur_batches=1) {
             for(int i=0; i<cur_batches; ++i){
 
                 cv::bitwise_and(segmented[i], masks[i], segmented[i]);
@@ -216,6 +218,7 @@ class SegmentationNN {
                 cv::imshow("segmented", segmented[i]);
                 cv::waitKey(1);
             }
+            return segmented[0];
         }
 
 };

include/tkDNN/kernelsThrust.h

@@ -2,6 +2,7 @@
 #define KERNELSTHRUST_H
 
 
+#include <thrust/extrema.h>
 #include <thrust/sort.h>
 #include <thrust/execution_policy.h>
 #include <thrust/functional.h>
@@ -9,6 +10,8 @@
 #include <thrust/iterator/constant_iterator.h>
 #include <thrust/gather.h>
 #include <thrust/copy.h>
+#include <thrust/device_ptr.h>
+
 
 #include "tkdnn.h"
 
@@ -36,4 +39,6 @@ void topKxyAddOffset(int * ids_begin, const int K, const int size, int *intxs_be
 void bboxes(int * ids_begin, const int K, const int size, float *xs_begin, float *ys_begin, 
             dnnType *src_begin, float *bbx0, float *bbx1, float *bby0, float *bby1, float *src_out, int *ids_out);
 
+void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w);
+
 #endif //KERNELSTHRUST_H

src/kernels/postprocessing.cu

@@ -34,6 +34,25 @@ void sortAndTopKonDevice(dnnType *src_begin, int *idsrc, float *topk_scores, int
     sortAndTopK_kernel<<<blocks, threads, 0>>>(src_begin, idsrc, topk_scores, topk_inds, topk_ys, topk_xs, size, K);   
 }
 
+__global__ 
+void maxElem_kernel(float *src_begin, float *dst_begin, const int n_classes, const int size){
+    int i = blockDim.x*blockIdx.x + threadIdx.x;
+    if (i > size)
+        return;
+
+    thrust::device_ptr<float> dPbeg ( &src_begin[i*n_classes] ) ;
+    thrust::device_ptr<float> dPend = dPbeg + n_classes;
+    thrust::device_ptr<float> result = thrust::max_element(thrust::device,dPbeg, dPend);
+
+    dst_begin[i] = result - dPbeg;
+}
+
+void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w){
+    int blocks = (h*w)/32+1;
+    int threads = 32;
+    maxElem_kernel<<<blocks, threads, 0>>>(src_begin, dst_begin, c, h*w);   
+}
+
 void topKxyclasses(int *ids_begin, int *ids_end, const int K, const int size, const int wh, int *clses, int *xs, int *ys){    
     thrust::transform(thrust::device, ids_begin, ids_end, thrust::make_constant_iterator(wh), clses, thrust::divides<int>());
     thrust::transform(thrust::device, ids_begin, ids_end, thrust::make_constant_iterator(wh), ids_begin, thrust::modulus<int>());