From bbcc33c0cfb64faac60f70a6c99f6193998b3735 Mon Sep 17 00:00:00 2001
From: Micaela Verucchi <micaelaverucchi@gmail.com>
Date: Fri, 20 Mar 2020 21:14:12 +0100
Subject: [PATCH] Refactoring for detection NN

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
---
 demo/demo/demo.cpp                 |  74 +++----
 demo/demo/map.cpp                  |   2 +-
 include/tkDNN/CenternetDetection.h | 180 ++++++----------
 include/tkDNN/DetectionNN.h        |  99 +++++++++
 include/tkDNN/MobilenetDetection.h |  71 ++-----
 include/tkDNN/Yolo3Detection.h     |  93 +++------
 include/tkDNN/utils.h              |   2 +-
 src/CenternetDetection.cpp         | 323 +++++++++++++----------------
 src/MobilenetDetection.cpp         | 240 +++++++++++----------
 src/Yolo3Detection.cpp             | 219 ++++++++++---------
 src/utils.cpp                      |  12 ++
 11 files changed, 637 insertions(+), 678 deletions(-)
 create mode 100644 include/tkDNN/DetectionNN.h

diff --git a/demo/demo/demo.cpp b/demo/demo/demo.cpp
index 90ff800..ccc5c7f 100644
--- a/demo/demo/demo.cpp
+++ b/demo/demo/demo.cpp
@@ -31,25 +31,34 @@ int main(int argc, char *argv[]) {
     char ntype = 'y';
     if(argc > 3)
         ntype = argv[3][0]; 
+    int n_classes = 80;
+    if(argc > 4)
+        n_classes = atoi(argv[4]); 
 
     tk::dnn::Yolo3Detection yolo;
     tk::dnn::CenternetDetection cnet;
-    tk::dnn::MobilenetDetection mbnet;    
+    tk::dnn::MobilenetDetection mbnet;  
+
+    tk::dnn::DetectionNN *detNN;  
+
     switch(ntype)
     {
         case 'y':
-            yolo.init(net);
+            detNN = &yolo;
             break;
         case 'c':
-            cnet.init(net);
+            detNN = &cnet;
             break;
         case 'm':
-            mbnet.init(net, 512, 81);
+            detNN = &mbnet;
+            n_classes++;
             break;
         default:
         FatalError("Network type not allowed (3rd parameter)\n");
     }
 
+    detNN->init(net, n_classes);
+
     gRun = true;
 
     cv::VideoCapture cap(input);
@@ -79,25 +88,11 @@ int main(int argc, char *argv[]) {
  
         // this will be resized to the net format
         dnn_input = frame.clone();
-        // TODO: async infer
-        switch(ntype)
-        {
-            case 'y':
-                yolo.update(dnn_input);
-                frame = yolo.draw(frame);
-                break;
-            case 'c':
-                cnet.update(dnn_input);
-                frame = cnet.draw(dnn_input);
-                break;
-            case 'm':
-                mbnet.update(dnn_input);
-                frame = mbnet.draw();
-                break;
-            default:
-                FatalError("Network type not allowed!\n");
-        }
-                
+        
+        //inference
+        detNN->update(dnn_input);
+        frame = detNN->draw(frame);
+
         cv::imshow("detection", frame);
         cv::waitKey(1);
         if(SAVE_RESULT)
@@ -106,32 +101,13 @@ int main(int argc, char *argv[]) {
 
     std::cout<<"detection end\n";   
     double mean = 0; 
-    switch(ntype)
-    {
-        case 'y':
-            std::cout<<COL_GREENB<<"\n\nTime stats:\n";
-            std::cout<<"Min: "<<*std::min_element(yolo.stats.begin(), yolo.stats.end())<<" ms\n";    
-            std::cout<<"Max: "<<*std::max_element(yolo.stats.begin(), yolo.stats.end())<<" ms\n";    
-            for(int i=0; i<yolo.stats.size(); i++) mean += yolo.stats[i]; mean /= yolo.stats.size();
-            std::cout<<"Avg: "<<mean<<" ms\n"<<COL_END;  
-            break;
-        case 'c':
-            std::cout<<COL_GREENB<<"\n\nTime stats:\n";
-            std::cout<<"Min: "<<*std::min_element(cnet.stats.begin(), cnet.stats.end())<<" ms\n";    
-            std::cout<<"Max: "<<*std::max_element(cnet.stats.begin(), cnet.stats.end())<<" ms\n";    
-            for(int i=0; i<cnet.stats.size(); i++) mean += cnet.stats[i]; mean /= cnet.stats.size();
-            std::cout<<"Avg: "<<mean<<" ms\n"<<COL_END;
-            break;
-        case 'm':
-            std::cout<<COL_GREENB<<"\n\nTime stats:\n";
-            std::cout<<"Min: "<<*std::min_element(mbnet.stats.begin(), mbnet.stats.end())<<" ms\n";    
-            std::cout<<"Max: "<<*std::max_element(mbnet.stats.begin(), mbnet.stats.end())<<" ms\n";    
-            for(int i=0; i<mbnet.stats.size(); i++) mean += mbnet.stats[i]; mean /= mbnet.stats.size();
-            std::cout<<"Avg: "<<mean<<" ms\n"<<COL_END;   
-            break;
-        default:
-            FatalError("Network type not allowed!\n");
-    }
+    
+    std::cout<<COL_GREENB<<"\n\nTime stats:\n";
+    std::cout<<"Min: "<<*std::min_element(detNN->stats.begin(), detNN->stats.end())<<" ms\n";    
+    std::cout<<"Max: "<<*std::max_element(detNN->stats.begin(), detNN->stats.end())<<" ms\n";    
+    for(int i=0; i<detNN->stats.size(); i++) mean += detNN->stats[i]; mean /= detNN->stats.size();
+    std::cout<<"Avg: "<<mean<<" ms\n"<<COL_END;   
+    
 
     return 0;
 }
diff --git a/demo/demo/map.cpp b/demo/demo/map.cpp
index aafa448..fcfd5a5 100644
--- a/demo/demo/map.cpp
+++ b/demo/demo/map.cpp
@@ -76,7 +76,7 @@ int main(int argc, char *argv[])
             cnet.init(net);
             break;
         case 'm':
-            mbnet.init(net, 512, 81);
+            mbnet.init(net, 81);
             break;
         default:
         FatalError("Network type not allowed (3rd parameter)\n");
diff --git a/include/tkDNN/CenternetDetection.h b/include/tkDNN/CenternetDetection.h
index e4bb358..4e987e3 100644
--- a/include/tkDNN/CenternetDetection.h
+++ b/include/tkDNN/CenternetDetection.h
@@ -1,134 +1,90 @@
-#include <iostream>
-#include <cstring>
-#include <signal.h>
-#include <stdlib.h>     /* srand, rand */
-#include <unistd.h>
-#include <mutex>
-#include "utils.h"
-#include <time.h>
+#ifndef CENTERNETDETECTION_H
+#define CENTERNETDETECTION_H
+
 #include "kernels.h"
+#include <opencv2/videoio.hpp>
+#include "opencv2/opencv.hpp"
+#include <time.h>
 #include <vector>
 #include <numeric>      // std::iota
 #include <algorithm>    // std::sort
 
+#include "DetectionNN.h"
 
-#include <opencv2/core/core.hpp>
-#include <opencv2/highgui/highgui.hpp>
-#include <opencv2/imgproc/imgproc.hpp>
-#include "opencv2/opencv.hpp"
-
-#include "tkdnn.h"
 #include "sorting.h"
 
-namespace tk { namespace dnn {
 
-/**
- * 
- * @author Francesco Gatti
- */
-class CenternetDetection {
+namespace tk { namespace dnn { 
 
-    private:
-        tk::dnn::NetworkRT *netRT = nullptr;
-        dnnType *input_d;
+class CenternetDetection : public DetectionNN
+{
+private:
+    std::vector<std::string> classesNames;
 
-        int ndets = 0;
-        // tk::dnn::Yolo::detection *dets = nullptr;
+    tk::dnn::dataDim_t dim;
+    tk::dnn::dataDim_t dim2;
+    tk::dnn::dataDim_t dim_hm; 
+    tk::dnn::dataDim_t dim_wh; 
+    tk::dnn::dataDim_t dim_reg;
+    float *topk_scores;
+    int *topk_inds_;
+    float *topk_ys_;
+    float *topk_xs_;
+    int *ids_d, *ids_, *ids_2, *ids_2d;
 
-        cv::Mat imageOrig;        
-        // std::vector< cv::cuda::GpuMat > bgr;
+    float *scores, *scores_d;
+    int *clses, *clses_d; 
+    int *topk_inds_d;
+    float *topk_ys_d; 
+    float *topk_xs_d;
+    int *inttopk_xs_d, *inttopk_ys_d;
 
-        // variable to test cnet on dog pictures
-        tk::dnn::dataDim_t dim;
-        tk::dnn::dataDim_t dim2;
-        cv::Size sz, sz_old;
-        const char *input_bin = "../tests/resnet101_cnet/debug/input.bin";
 
-        cv::cuda::Stream stream;
-        struct threshold op;
-        // pre-process
-        tk::dnn::dataDim_t dim_hm; 
-        tk::dnn::dataDim_t dim_wh; 
-        tk::dnn::dataDim_t dim_reg;
-        float *topk_scores;
-        int *topk_inds_;
-        float *topk_ys_;
-        float *topk_xs_;
-        int *ids_d, *ids_, *ids_2, *ids_2d;
-
-        float *scores, *scores_d;
-        int *clses, *clses_d; 
-        int *topk_inds_d;
-        float *topk_ys_d; 
-        float *topk_xs_d;
-        int *inttopk_xs_d, *inttopk_ys_d;
+    float *bbx0, *bby0, *bbx1, *bby1;
+    float *bbx0_d, *bby0_d, *bbx1_d, *bby1_d;
     
+    float *target_coords;
 
-        float *bbx0, *bby0, *bbx1, *bby1;
-        float *bbx0_d, *bby0_d, *bbx1_d, *bby1_d;
-        
-        float *target_coords;
-        
-        
+    #ifdef OPENCV_CUDA
+        float *mean_d;
+        float *stddev_d;
+    #else
+        cv::Vec<float, 3> mean;
+        cv::Vec<float, 3> stddev;
+        dnnType *input;
+    #endif
 
-        #ifdef OPENCV_CUDA
-            float *mean_d;
-            float *stddev_d;
-        #else
-            cv::Vec<float, 3> mean;
-            cv::Vec<float, 3> stddev;
-            dnnType *input;
-        #endif
+    float *d_ptrs;
 
-        float *d_ptrs;
-        
-        cv::Mat src;
-        cv::Mat dst;
-        cv::Mat dst2;  
-        cv::Mat trans, trans2;
-        //processing
-        float toll = 0.000001;
-        int K = 100;
-        int width = 128;//56;        // TODO
+    cv::Mat src;
+    cv::Mat dst;
+    cv::Mat dst2;  
+    cv::Mat trans, trans2;
+    //processing
+    float toll = 0.000001;
+    int K = 100;
+    int width = 128;//56;        // TODO
+
+    // pointer used in the kernels
+    float *src_out;
+    int *ids_out;
+
+    struct threshold op;
     
-        // pointer used in the kernels
-        float *src_out;
-        int *ids_out;
-
-        void preprocess();
-    public:
-        dnnType *rt_out[4];
-
-        float inp_height = 512;//224;//512;
-        float inp_width = 512;//224;//512;
-    
-        int classes = 80;
-        int num = 0;
-        int n_masks = 0;
-        float thresh = 0.3;
-        cv::Scalar colors[256];
-
-        // this is filled with results
-        std::vector<tk::dnn::box> detected;
-        // draw
-        std::vector<std::string> coco_class_name;
-
-        // keep track of inference times (ms)
-        std::vector<double> stats;
-        
-        CenternetDetection() {}
-
-        virtual ~CenternetDetection() {}
-
-        /**
-         * Method used for inizialize the class
-         * 
-         * @return Success of the initialization
-         */
-        bool init(std::string tensor_path);
-        cv::Mat draw(cv::Mat &frame);
-        void update(cv::Mat &frame);
+public:
+    CenternetDetection() {};
+    ~CenternetDetection() {}; 
 
+    bool init(const std::string& tensor_path, const int n_classes=80);
+    void preprocess(cv::Mat &frame);
+    void update(cv::Mat &frame);
+    void postprocess(dnnType **rt_out, const int n_out);
+    cv::Mat draw(cv::Mat &frame);
 };
 
-}}
+
+} // namespace dnn
+} // namespace tk
+
+
+#endif /*CENTERNETDETECTION_H*/
\ No newline at end of file
diff --git a/include/tkDNN/DetectionNN.h b/include/tkDNN/DetectionNN.h
new file mode 100644
index 0000000..ddc2d43
--- /dev/null
+++ b/include/tkDNN/DetectionNN.h
@@ -0,0 +1,99 @@
+#ifndef DETECTIONNN_H
+#define DETECTIONNN_H
+
+#include <iostream>
+#include <signal.h>
+#include <stdlib.h>     /* srand, rand */
+#include <unistd.h>
+#include <mutex>
+#include "utils.h"
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include "tkdnn.h"
+
+// #define OPENCV_CUDA //if OPENCV has been compiled with CUDA and contrib.
+
+namespace tk { namespace dnn {
+
+enum networkType_t{
+    NETWORK_YOLO3,
+    NETWORK_MOBILENETSSDLITE,
+    NETWORK_CENTERNET
+};
+
+class DetectionNN {
+
+    protected:
+        tk::dnn::NetworkRT *netRT = nullptr;
+        dnnType *input_d;
+
+        cv::Size originalSize;
+
+        cv::Scalar colors[256];
+
+#ifdef OPENCV_CUDA
+        cv::cuda::GpuMat bgr[3];
+        cv::cuda::GpuMat imagePreproc;
+#else
+        cv::Mat bgr[3];
+        cv::Mat imagePreproc;
+        dnnType *input;
+#endif
+
+    public:
+        int classes = 0;
+        float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
+
+        std::vector<tk::dnn::box> detected; /*bounding boxes in output*/
+        std::vector<double> stats; /*keeps track of inference times (ms)*/
+
+        DetectionNN() {};
+        ~DetectionNN(){};
+
+        /**
+         * Method used to inialize the class, allocate memory and compute 
+         * needed data.
+         * 
+         * @param path to the rt file og the NN.
+         * @return true if everything is correct, false otherwise.
+         */
+        virtual bool init(const std::string& tensor_path, const int n_classes=80) = 0;
+        
+        /**
+         * This method preprocess the image, before feeding it to the NN.
+         *
+         * @param original frame to adapt for inference.
+         */
+        virtual void preprocess(cv::Mat &frame) = 0;
+
+        /**
+         * This method performs the inference of the NN.
+         * 
+         * @param frame to run inference on.
+         */
+        virtual void update(cv::Mat &frame) = 0;
+
+        /**
+         * This method postprocess the output of the NN to obtain the correct 
+         * boundig boxes. 
+         * 
+         * @param outputs of the inference
+         * @param number of outputs of the inference
+         */
+        virtual void postprocess(dnnType **rt_out, const int n_out) = 0;
+
+        /**
+         * Method to draw boundixg boxes and labels on a frame.
+         * 
+         * @param orginal frame to draw bounding box on.
+         * @return frame with boundig boxes.
+         */
+        virtual cv::Mat draw(cv::Mat &frame) = 0;
+};
+
+}}
+
+#endif /* DETECTIONNN_H*/
diff --git a/include/tkDNN/MobilenetDetection.h b/include/tkDNN/MobilenetDetection.h
index f4bec95..44fd120 100644
--- a/include/tkDNN/MobilenetDetection.h
+++ b/include/tkDNN/MobilenetDetection.h
@@ -1,18 +1,15 @@
 #ifndef MOBILENETDETECTION_H
 #define MOBILENETDETECTION_H
 
-#include <iostream>
-
-#include <opencv2/core/core.hpp>
-#include <opencv2/highgui/highgui.hpp>
 #include <opencv2/videoio.hpp>
-#include <opencv2/imgproc/imgproc.hpp>
 #include "opencv2/opencv.hpp"
 
-#include "tkdnn.h"
+#include "DetectionNN.h"
 
 #define N_COORDS 4
+#define N_SSDSPEC 6
 
+namespace tk { namespace dnn { 
 
 struct SSDSpec
 {
@@ -24,10 +21,9 @@ struct SSDSpec
     int ratio2 = 0;
 
     SSDSpec() {}
-
-    SSDSpec(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2) : featureSize(feature_size), shrinkage(shrinkage), boxWidth(box_width), boxHeight(box_height),
-                                                                                                      ratio1(ratio1), ratio2(ratio2) {}
-
+    SSDSpec(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2) : 
+                    featureSize(feature_size), shrinkage(shrinkage), boxWidth(box_width), 
+                    boxHeight(box_height), ratio1(ratio1), ratio2(ratio2) {}
     void setAll(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2)
     {
         this->featureSize = feature_size;
@@ -37,74 +33,47 @@ struct SSDSpec
         this->ratio1 = ratio1;
         this->ratio2 = ratio2;
     }
-
     void print()
     {
-        std::cout << "fsize: " << featureSize << "\tshrinkage: " << shrinkage << "\t box W:" << boxWidth << "\tbox H: " << boxHeight << "\t x ratio:" << ratio1 << "\t y ratio:" << ratio2 << std::endl;
+        std::cout << "fsize: " << featureSize << "\tshrinkage: " << shrinkage << 
+                    "\t box W:" << boxWidth << "\tbox H: " << boxHeight << 
+                    "\t x ratio:" << ratio1 << "\t y ratio:" << ratio2 << std::endl;
     }
 };
 
-
-namespace tk
+class MobilenetDetection : public DetectionNN
 {
-namespace dnn
-{
-class MobilenetDetection
-{
-
 private:
-    tk::dnn::NetworkRT *netRT = nullptr;
-
-    int classes;
     float IoUThreshold = 0.45;
     float centerVariance = 0.1;
     float sizeVariance = 0.2;
-    float confThreshold = 0.4;
     int imageSize;
 
     float *priors = nullptr;
     int nPriors = 0;
-
-    cv::Mat origImg;
-
-    float *input, *input_d;
     float *locations_h, *confidences_h;
 
-    tk::dnn::dataDim_t dim;
-
-    dnnType *conf;
-    dnnType *loc;
-
-    float __colors[6][3] = {{1, 0, 1}, {0, 0, 1}, {0, 1, 1}, {0, 1, 0}, {1, 1, 0}, {1, 0, 0}};
-    int baseline = 0;
-    float fontScale = 0.5;
-    int thickness = 2;
+    std::vector<std::string> classesNames;
 
     void generate_ssd_priors(const SSDSpec *specs, const int n_specs, bool clamp = true);
     void convert_locatios_to_boxes_and_center();
     float iou(const tk::dnn::box &a, const tk::dnn::box &b);
-    void preprocess();
-    std::vector<tk::dnn::box> postprocess(const int width, const int height);
-    float get_color2(int c, int x, int max);
 
-    cv::Scalar colors[256];
-    std::vector<std::string> classesNames;
+    
 
 public:
-    // keep track of inference times (ms)
-    std::vector<double> stats;
-    std::vector<tk::dnn::box> detected;
+    MobilenetDetection() {};
+    ~MobilenetDetection() {}; 
 
-    MobilenetDetection() {}
-    ~MobilenetDetection() {}
-
-    void init(std::string tensor_path, int input_size, int n_classes);
-    cv::Mat draw();
-    void update(cv::Mat &img);
+    bool init(const std::string& tensor_path, const int n_classes);
+    void preprocess(cv::Mat &frame);
+    void update(cv::Mat &frame);
+    void postprocess(dnnType **rt_out, const int n_out);
+    cv::Mat draw(cv::Mat &frame);
 };
 
+
 } // namespace dnn
 } // namespace tk
 
-
 #endif /*MOBILENETDETECTION_H*/
\ No newline at end of file
diff --git a/include/tkDNN/Yolo3Detection.h b/include/tkDNN/Yolo3Detection.h
index 80d179e..148a53f 100644
--- a/include/tkDNN/Yolo3Detection.h
+++ b/include/tkDNN/Yolo3Detection.h
@@ -1,73 +1,36 @@
-#ifndef YOLODETECTION_H
-#define YOLODETECTION_H
+#ifndef Yolo3Detection_H
+#define Yolo3Detection_H
+#include <opencv2/videoio.hpp>
+#include "opencv2/opencv.hpp"
 
-#include <iostream>
-#include <signal.h>
-#include <stdlib.h>     /* srand, rand */
-#include <unistd.h>
-#include <mutex>
-#include "utils.h"
+#include "DetectionNN.h"
 
-#include <opencv2/core/core.hpp>
-#include <opencv2/highgui/highgui.hpp>
-#include <opencv2/imgproc/imgproc.hpp>
+namespace tk { namespace dnn { 
 
-#include "tkdnn.h"
+class Yolo3Detection : public DetectionNN
+{
+private:
+    int num = 0;
+    int nMasks = 0;
+    int nDets = 0;
+    tk::dnn::Yolo::detection *dets = nullptr;
+    tk::dnn::Yolo* yolo[3];
 
-namespace tk { namespace dnn {
-
-/**
- * 
- * @author Francesco Gatti
- */
-class Yolo3Detection {
-
-    private:
-        tk::dnn::NetworkRT *netRT = nullptr;
-        tk::dnn::Yolo* yolo[3];
-        dnnType *input, *input_d;
-
-        int ndets = 0;
-        tk::dnn::Yolo::detection *dets = nullptr;
-
-        cv::Mat imageF;
-        cv::Mat bgr[3]; 
-
-    public:
-        int classes = 0;
-        int num = 0;
-        int n_masks = 0;
-        float thresh = 0.3;
-        cv::Scalar colors[256];
-
-        // this is filled with results
-        std::vector<tk::dnn::box> detected;
-
-        // keep track of inference times (ms)
-        std::vector<double> stats;
-
-        Yolo3Detection() {}
-
-        virtual ~Yolo3Detection() {}
-
-        /**
-         * Method used for inizialize the class
-         * 
-         * @return Success of the initialization
-         */
-        bool init(std::string tensor_path);
-        cv::Mat draw(cv::Mat &frame);
-        void update(cv::Mat &frame);
-
-        tk::dnn::Yolo* getYoloLayer(int n=0) {
-            if(n<3)
-                return yolo[n];
-            else 
-                return nullptr;
-        }
+    tk::dnn::Yolo* getYoloLayer(int n=0);
+    
+public:
+    Yolo3Detection() {};
+    ~Yolo3Detection() {}; 
 
+    bool init(const std::string& tensor_path, const int n_classes=80);
+    void preprocess(cv::Mat &frame);
+    void update(cv::Mat &frame);
+    void postprocess(dnnType **rt_out, const int n_out);
+    cv::Mat draw(cv::Mat &frame);
 };
 
-}}
 
-#endif /* YOLODETECTION_H*/
+} // namespace dnn
+} // namespace tk
+
+#endif /* Yolo3Detection_H*/
diff --git a/include/tkDNN/utils.h b/include/tkDNN/utils.h
index 1a6df95..2b8ba8a 100644
--- a/include/tkDNN/utils.h
+++ b/include/tkDNN/utils.h
@@ -14,7 +14,6 @@
 
 #define dnnType float
 
-#define OPENCV_CUDA
 
 // Colored output
 #define COL_END "\033[0m"
@@ -96,6 +95,7 @@ void downloadWeightsifDoNotExist(const std::string& input_bin, const std::string
 void readBinaryFile(std::string fname, int size, dnnType** data_h, dnnType** data_d, int seek = 0, bool skipLoad = false);
 int checkResult(int size, dnnType *data_d, dnnType *correct_d, bool device = true);
 void printDeviceVector(int size, dnnType* vec_d, bool device = true);
+float getColor(const int c, const int x, const int max);
 void resize(int size, dnnType **data);
 
 void matrixTranspose(cublasHandle_t handle, dnnType* srcData, dnnType* dstData, int rows, int cols);
diff --git a/src/CenternetDetection.cpp b/src/CenternetDetection.cpp
index b1e8ff9..29cdec6 100644
--- a/src/CenternetDetection.cpp
+++ b/src/CenternetDetection.cpp
@@ -1,33 +1,19 @@
-#ifndef CENTERNETDETECTION_H
-#define CENTERNETDETECTION_H
+#include "CenternetDetection.h"
 
 #include "CenternetDetection.h"
-#include "opencv2/imgproc/imgproc.hpp"
-// #include <opencv2/cudawarping.hpp>
-// #include <opencv2/cudaarithm.hpp>
-
-
 
 namespace tk { namespace dnn {
 
-float __colors[6][3] = { {1,0,1}, {0,0,1},{0,1,1},{0,1,0},{1,1,0},{1,0,0} };
-float get_color2(int c, int x, int max)
-{
-    float ratio = ((float)x/max)*5;
-    int i = floor(ratio);
-    int j = ceil(ratio);
-    ratio -= i;
-    float r = (1-ratio) * __colors[i % 6][c % 3] + ratio*__colors[j % 6][c % 3];
-    //printf("%f\n", r);
-    return r;
-}
 
-bool CenternetDetection::init(std::string tensor_path) {
+bool CenternetDetection::init(const std::string& tensor_path, const int n_classes)
+{
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
-    
-    dim = tk::dnn::dataDim_t(1, 3, 512, 512, 1);
-    const char *coco_class_name_[] = {
+    classes = n_classes;
+
+    dim = netRT->input_dim;
+
+    const char *coco_class_name[] = {
             "person", "bicycle", "car", "motorcycle", "airplane", 
             "bus", "train", "truck", "boat", "traffic light", "fire hydrant",
             "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
@@ -42,19 +28,24 @@ bool CenternetDetection::init(std::string tensor_path) {
             "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
             "scissors", "teddy bear", "hair drier", "toothbrush"
             };
-    coco_class_name = std::vector<std::string>(coco_class_name_, std::end( coco_class_name_ ));
+    classesNames = std::vector<std::string>(coco_class_name, std::end( coco_class_name));
+    
+    for(int c=0; c<classes; c++) {
+        int offset = c*123457 % classes;
+        float r = getColor(2, offset, classes);
+        float g = getColor(1, offset, classes);
+        float b = getColor(0, offset, classes);
+        colors[c] = cv::Scalar(int(255.0*b), int(255.0*g), int(255.0*r));
+    }
+    
     src = cv::Mat(cv::Size(2,3), CV_32F);
     dst = cv::Mat(cv::Size(2,3), CV_32F);
     dst2 = cv::Mat(cv::Size(2,3), CV_32F);
     trans = cv::Mat(cv::Size(3,2), CV_32F);
     trans2 = cv::Mat(cv::Size(3,2), CV_32F);
-    // dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
 
     checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot()));
 
-    // dim_hm = tk::dnn::dataDim_t(1, 80, 56, 56, 1);
-    // dim_wh = tk::dnn::dataDim_t(1, 2, 56, 56, 1);
-    // dim_reg = tk::dnn::dataDim_t(1, 2, 56, 56, 1);
     dim_hm = tk::dnn::dataDim_t(1, 80, 128, 128, 1);
     dim_wh = tk::dnn::dataDim_t(1, 2, 128, 128, 1);
     dim_reg = tk::dnn::dataDim_t(1, 2, 128, 128, 1);
@@ -79,19 +70,16 @@ bool CenternetDetection::init(std::string tensor_path) {
 
     checkCuda( cudaMallocHost(&scores, K *sizeof(float)) );
     checkCuda( cudaMalloc(&scores_d, K *sizeof(float)) );
+
     checkCuda( cudaMallocHost(&clses, K *sizeof(int)) );
     checkCuda( cudaMalloc(&clses_d, K *sizeof(int)) );
-    // checkCuda( cudaMallocHost(&topk_inds, K *sizeof(int)) );
+
     checkCuda( cudaMalloc(&topk_inds_d, K *sizeof(int)) );
     checkCuda( cudaMalloc(&topk_ys_d, K *sizeof(float)) );     
     checkCuda( cudaMalloc(&topk_xs_d, K *sizeof(float)) ); 
-    // checkCuda( cudaMalloc(&intid, K *sizeof(int)) );
     checkCuda( cudaMalloc(&inttopk_ys_d, K *sizeof(int)) );
     checkCuda( cudaMalloc(&inttopk_xs_d, K *sizeof(int)) );
-    
-    // checkCuda( cudaMalloc(&ids_d, dim_hm.c * K*sizeof(int)) );
 
-    // checkCuda( cudaMallocHost(&wh_aus, dim_wh.tot()*sizeof(dnnType)) ); 
     checkCuda( cudaMallocHost(&bbx0, K * sizeof(float)) ); 
     checkCuda( cudaMallocHost(&bby0, K * sizeof(float)) ); 
     checkCuda( cudaMallocHost(&bbx1, K * sizeof(float)) ); 
@@ -119,14 +107,11 @@ bool CenternetDetection::init(std::string tensor_path) {
 #endif
 
     checkCuda( cudaMalloc(&d_ptrs, dim.c * dim.h*dim.w * sizeof(float)) );
-    // mean << 0.408, 0.447, 0.47;
-    // stddev << 0.289, 0.274, 0.278;
 
     // Alloc array used in the kernel 
     checkCuda( cudaMalloc(&src_out, K *sizeof(float)) );
     checkCuda( cudaMalloc(&ids_out, K *sizeof(int)) );
-    // checkCuda( cudaFree(src_out) );
-    // checkCuda( cudaFree(ids_out) );
+
     dst2.at<float>(0,0)=width * 0.5;
     dst2.at<float>(0,1)=width * 0.5;
     dst2.at<float>(1,0)=width * 0.5;
@@ -137,60 +122,17 @@ bool CenternetDetection::init(std::string tensor_path) {
     
 }
 
-cv::Mat CenternetDetection::draw(cv::Mat &imageOrig) {
 
-    tk::dnn::box b;
-    int x0, w, x1, y0, h, y1;
-    int objClass;
-    std::string det_class;
-    int baseline = 0;
-    float fontScale = 0.5;
-    int thickness = 2;   
-
-    for(int c=0; c<classes; c++) {
-        int offset = c*123457 % classes;
-        float r = get_color2(2, offset, classes);
-        float g = get_color2(1, offset, classes);
-        float b = get_color2(0, offset, classes);
-        colors[c] = cv::Scalar(int(255.0*b), int(255.0*g), int(255.0*r));
-    }
-    int num_detected = detected.size();
-    for (int i = 0; i < num_detected; i++){
-        b = detected[i];
-        x0 = b.x;
-        w = b.w;
-        x1 = b.x + w;
-        y0 = b.y;
-        h = b.h;
-        y1 = b.y + h;
-        objClass = b.cl;
-        det_class = coco_class_name[objClass];
-        cv::rectangle(imageOrig, cv::Point(x0, y0), cv::Point(x1, y1), colors[objClass], 2);
-        // draw label
-        cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
-        cv::rectangle(imageOrig, cv::Point(x0, y0), cv::Point((x0 + textSize.width - 2), (y0 - textSize.height - 2)), colors[b.cl], -1);
-        cv::putText(imageOrig, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
-        
-    }
-    return imageOrig;
-    // cv::namedWindow("cnet", cv::WINDOW_NORMAL);
-    // cv::imshow("cnet", imageOrig);
-    // cv::waitKey(10000);
-}
-
-
-void CenternetDetection::preprocess()
+void CenternetDetection::preprocess(cv::Mat &frame)
 {
-
-    auto start_t = std::chrono::steady_clock::now();
-    auto step_t = std::chrono::steady_clock::now();
-    auto end_t = std::chrono::steady_clock::now();
-
      // -----------------------------------pre-process ------------------------------------------
-    // it will resize the images to `224 x 224` in GETTING_STARTED.md
-    cv::Size sz = imageOrig.size();
-    std::cout<<"image: "<<sz.width<<", "<<sz.height<<std::endl;
     
+    // auto start_t = std::chrono::steady_clock::now();
+    // auto step_t = std::chrono::steady_clock::now();
+    // auto end_t = std::chrono::steady_clock::now();
+    cv::Size sz = originalSize;
+    std::cout<<"image: "<<sz.width<<", "<<sz.height<<std::endl;
+    cv::Size sz_old;
     float scale = 1.0;
     float new_height = sz.height * scale;
     float new_width = sz.width * scale;
@@ -214,10 +156,10 @@ void CenternetDetection::preprocess()
         src.at<float>(0,1)=c[1];
         src.at<float>(1,0)=c[0];
         src.at<float>(1,1)=c[1] + s[0] * -0.5;
-        dst.at<float>(0,0)=inp_width * 0.5;
-        dst.at<float>(0,1)=inp_height * 0.5;
-        dst.at<float>(1,0)=inp_width * 0.5;
-        dst.at<float>(1,1)=inp_height * 0.5 +  inp_width * -0.5; 
+        dst.at<float>(0,0)=netRT->input_dim.w * 0.5;
+        dst.at<float>(0,1)=netRT->input_dim.h * 0.5;
+        dst.at<float>(1,0)=netRT->input_dim.w * 0.5;
+        dst.at<float>(1,1)=netRT->input_dim.h * 0.5 +  netRT->input_dim.w * -0.5; 
         
         src.at<float>(2,0)=src.at<float>(1,0) + (-src.at<float>(0,1)+src.at<float>(1,1) );
         src.at<float>(2,1)=src.at<float>(1,1) + (src.at<float>(0,0)-src.at<float>(1,0) );
@@ -225,91 +167,87 @@ void CenternetDetection::preprocess()
         dst.at<float>(2,1)=dst.at<float>(1,1) + (dst.at<float>(0,0)-dst.at<float>(1,0) );
 
         trans = cv::getAffineTransform( src, dst );
-        end_t = std::chrono::steady_clock::now();
-        std::cout << " TIME gett affine trans: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-        step_t = end_t;
+        // end_t = std::chrono::steady_clock::now();
+        // std::cout << " TIME gett affine trans: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+        // step_t = end_t;
 
         trans2 = cv::getAffineTransform( dst2, src );
-
-        end_t = std::chrono::steady_clock::now();
-        std::cout << " TIME getAffineTrans 2: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-        step_t = end_t;
+        // end_t = std::chrono::steady_clock::now();
+        // std::cout << " TIME getAffineTrans 2: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+        // step_t = end_t;
     }
     sz_old = sz;
 #ifdef OPENCV_CUDA
     cv::cuda::GpuMat im_Orig; 
     cv::cuda::GpuMat imageF1_d, imageF2_d;
         
-    im_Orig = cv::cuda::GpuMat(imageOrig);
+    im_Orig = cv::cuda::GpuMat(frame);
     cv::cuda::resize (im_Orig, imageF1_d, cv::Size(new_width, new_height)); 
     checkCuda( cudaDeviceSynchronize() );
     
     sz = imageF1_d.size();
     std::cout<<"size: "<<sz.height<<" "<<sz.width<<" - "<<std::endl;
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME resize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME resize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
     
-    cv::cuda::warpAffine(imageF1_d, imageF2_d, trans, cv::Size(inp_width, inp_height), cv::INTER_LINEAR );
+    cv::cuda::warpAffine(imageF1_d, imageF2_d, trans, cv::Size(netRT->input_dim.w, netRT->input_dim.h), cv::INTER_LINEAR );
     checkCuda( cudaDeviceSynchronize() );
     
-    
     imageF2_d.convertTo(imageF1_d, CV_32FC3, 1/255.0); 
     checkCuda( cudaDeviceSynchronize() );
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME convert: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME convert: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
     
     dim2 = dim;
     cv::cuda::GpuMat bgr[3]; 
     cv::cuda::split(imageF1_d,bgr);//split source
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME split: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME split: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     for(int i=0; i<dim.c; i++)
         checkCuda( cudaMemcpy(d_ptrs + i*dim.h * dim.w, (float*)bgr[i].data, dim.h * dim.w * sizeof(float), cudaMemcpyDeviceToDevice) );
         
     normalize(d_ptrs, dim.c, dim.h, dim.w, mean_d, stddev_d);
     
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME normalize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME normalize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     checkCuda(cudaMemcpy(input_d, d_ptrs, dim2.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice));
     
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME Memcpy to input_d: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME Memcpy to input_d: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 #else
  
     cv::Mat imageF;
-    resize(imageOrig, imageF, cv::Size(new_width, new_height));
+    resize(frame, imageF, cv::Size(new_width, new_height));
     sz = imageF.size();
     std::cout<<"size: "<<sz.height<<" "<<sz.width<<" - "<<std::endl;
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME resize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME resize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     cv::Mat trans = cv::getAffineTransform( src, dst );
-    cv::warpAffine(imageF, imageF, trans, cv::Size(inp_width, inp_height), cv::INTER_LINEAR );
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME warpAffine: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    cv::warpAffine(imageF, imageF, trans, cv::Size(netRT->input_dim.w, netRT->input_dim.h), cv::INTER_LINEAR );
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME warpAffine: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     sz = imageF.size();
     std::cout<<"size: "<<sz.height<<" "<<sz.width<<" - "<<std::endl;
     imageF.convertTo(imageF, CV_32FC3, 1/255.0); 
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME convertto: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME convertto: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     dim2 = dim;
-
     //split channels
     cv::Mat bgr[3]; 
     cv::split(imageF,bgr);//split source
-
     for(int i=0; i<3; i++){
         bgr[i] = bgr[i] - mean[i];
         bgr[i] = bgr[i] / stddev[i];
@@ -322,28 +260,20 @@ void CenternetDetection::preprocess()
         // std::cout<<"i: "<<i<<", idx: "<<idx<<", ch: "<<ch<<std::endl;
         memcpy((void*)&input[idx], (void*)bgr[ch].data, imageF.rows*imageF.cols*sizeof(dnnType));
     }
-
     checkCuda(cudaMemcpyAsync(input_d, input, dim2.tot()*sizeof(dnnType), cudaMemcpyHostToDevice));
-
-
-
 #endif
-
 }
 
-void CenternetDetection::update(cv::Mat &image_orig) {
-    
-    imageOrig = image_orig;
-    if(!imageOrig.data) {
+void CenternetDetection::update(cv::Mat &frame) 
+{
+    originalSize = frame.size();
+    if(!frame.data) {
         std::cout<<"CENTERNET: NO IMAGE DATA\n";
         return;
     }  
     TIMER_START
-    auto start_t = std::chrono::steady_clock::now();
-    auto step_t = std::chrono::steady_clock::now();
-    auto end_t = std::chrono::steady_clock::now();
    
-    preprocess();
+    preprocess(frame);
 
     printCenteredTitle(" TENSORRT inference ", '=', 30); {
         dim2.print();
@@ -352,22 +282,33 @@ void CenternetDetection::update(cv::Mat &image_orig) {
         TIMER_STOP
         dim2.print();
     }
-    step_t = std::chrono::steady_clock::now();
-    
-    // ------------------------------------ process --------------------------------------------
-    
+
+    dnnType *rt_out[4];
     rt_out[0] = (dnnType *)netRT->buffersRT[1];
     rt_out[1] = (dnnType *)netRT->buffersRT[2];
     rt_out[2] = (dnnType *)netRT->buffersRT[3]; 
     rt_out[3] = (dnnType *)netRT->buffersRT[4]; 
+
+    postprocess(rt_out, 4);
+
+    // std::cout<<"TOTAL: \n";
+    TIMER_STOP
+    stats.push_back(t_ns);
+}
+
+void CenternetDetection::postprocess(dnnType **rt_out, const int n_out)
+{
+    // auto start_t = std::chrono::steady_clock::now();
+    // auto step_t = std::chrono::steady_clock::now();
+    // auto end_t = std::chrono::steady_clock::now();
+    // ------------------------------------ process --------------------------------------------
     activationSIGMOIDForward(rt_out[0], rt_out[0], dim_hm.tot());
     checkCuda( cudaDeviceSynchronize() );
 
     subtractWithThreshold(rt_out[0], rt_out[0] + dim_hm.tot(), rt_out[1], rt_out[0], op);
-
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME threshold: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME threshold: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
     // ----------- nms end
     // ----------- topk  
       
@@ -378,28 +319,25 @@ void CenternetDetection::update(cv::Mat &image_orig) {
     
     checkCuda( cudaMemcpy(ids_d, ids_, dim_hm.c * dim_hm.h * dim_hm.w*sizeof(int), cudaMemcpyHostToDevice) );    
 
-    sort(rt_out[0],
-            rt_out[0]+dim_hm.tot(),
-            ids_d);
+    sort(rt_out[0],rt_out[0]+dim_hm.tot(),ids_d);
     checkCuda( cudaDeviceSynchronize() );   
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME sort: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME sort: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
-    topk(rt_out[0], ids_d, K, scores_d,
-        topk_inds_d, topk_ys_d, topk_xs_d);
+    topk(rt_out[0], ids_d, K, scores_d, topk_inds_d, topk_ys_d, topk_xs_d);
     checkCuda( cudaDeviceSynchronize() );    
 
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME topk: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME topk: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
         
     checkCuda( cudaMemcpy(scores, scores_d, K *sizeof(float), cudaMemcpyDeviceToHost) );
 
     topKxyclasses(topk_inds_d, topk_inds_d+K, K, width, dim_hm.w*dim_hm.h, clses_d, inttopk_xs_d, inttopk_ys_d);
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME topk x y clses 2: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME topk x y clses 2: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     checkCuda( cudaMemcpy(topk_xs_d, (float *)inttopk_xs_d, K*sizeof(float), cudaMemcpyDeviceToDevice) );
     checkCuda( cudaMemcpy(topk_ys_d, (float *)inttopk_ys_d, K*sizeof(float), cudaMemcpyDeviceToDevice) );
@@ -411,9 +349,9 @@ void CenternetDetection::update(cv::Mat &image_orig) {
     topKxyAddOffset(topk_inds_d, K, dim_reg.h*dim_reg.w, inttopk_xs_d, inttopk_ys_d, topk_xs_d, topk_ys_d, rt_out[3], src_out, ids_out);
     // checkCuda( cudaDeviceSynchronize() );
     
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME add offset: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME add offset: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
     
     bboxes(topk_inds_d, K, dim_wh.h*dim_wh.w, topk_xs_d, topk_ys_d, rt_out[2], bbx0_d, bbx1_d, bby0_d, bby1_d, src_out, ids_out);
     // checkCuda( cudaDeviceSynchronize() );
@@ -423,9 +361,9 @@ void CenternetDetection::update(cv::Mat &image_orig) {
     checkCuda( cudaMemcpy(bbx1, bbx1_d, K * sizeof(float), cudaMemcpyDeviceToHost) ); 
     checkCuda( cudaMemcpy(bby1, bby1_d, K * sizeof(float), cudaMemcpyDeviceToHost) ); 
     
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME bboxes: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME bboxes: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 
     // ---------------------------------- post-process -----------------------------------------
     
@@ -459,7 +397,7 @@ void CenternetDetection::update(cv::Mat &image_orig) {
     for(int i = 0; i<classes; i++){
         for(int j=0; j<K; j++)
             if(clses[j] == i){
-                if(scores[j] > thresh){
+                if(scores[j] > confThreshold){
                     // std::cout<<"th: "<<scores[j]<<" - cl: "<<clses[j]<<" i: "<<i<<std::endl;
                     //add coco bbox
                     //det[0:4], i, det[4]
@@ -482,14 +420,43 @@ void CenternetDetection::update(cv::Mat &image_orig) {
             }
     }
     
-    end_t = std::chrono::steady_clock::now();
-    std::cout << " TIME detections: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
-    step_t = end_t;
-
-    std::cout<<"TOTAL: \n";
-    TIMER_STOP
-    stats.push_back(t_ns);
+    // end_t = std::chrono::steady_clock::now();
+    // std::cout << " TIME detections: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << "  us" << std::endl;
+    // step_t = end_t;
 }
+
+cv::Mat CenternetDetection::draw(cv::Mat &frame)
+{
+    tk::dnn::box b;
+    int x0, w, x1, y0, h, y1;
+    int objClass;
+    std::string det_class;
+    int baseline = 0;
+    float font_scale = 0.5;
+    int thickness = 2;   
+
+    int num_detected = detected.size();
+    for (int i = 0; i < num_detected; i++){
+        b = detected[i];
+        x0 = b.x;
+        w = b.w;
+        x1 = b.x + w;
+        y0 = b.y;
+        h = b.h;
+        y1 = b.y + h;
+        objClass = b.cl;
+        det_class = classesNames[objClass];
+        cv::rectangle(frame, cv::Point(x0, y0), cv::Point(x1, y1), colors[objClass], 2);
+        // draw label
+        cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
+        cv::rectangle(frame, cv::Point(x0, y0), cv::Point((x0 + textSize.width - 2), (y0 - textSize.height - 2)), colors[b.cl], -1);
+        cv::putText(frame, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
+        
+    }
+    return frame;
+}
+
+
 }}
 
-#endif /*CENTERNETDETECTION_H*/
+
diff --git a/src/MobilenetDetection.cpp b/src/MobilenetDetection.cpp
index b992fa8..adee317 100644
--- a/src/MobilenetDetection.cpp
+++ b/src/MobilenetDetection.cpp
@@ -4,14 +4,12 @@ bool boxProbCmp(const tk::dnn::box &a, const tk::dnn::box &b){
     return (a.prob > b.prob);
 }
 
-namespace tk{
-namespace dnn{
+namespace tk{ namespace dnn{
 
 void MobilenetDetection::generate_ssd_priors(const SSDSpec *specs, const int n_specs, bool clamp)
 {
     nPriors = 0;
-    for (int i = 0; i < n_specs; i++)
-    {
+    for (int i = 0; i < n_specs; i++){
         nPriors += specs[i].featureSize * specs[i].featureSize * 6;
     }
 
@@ -131,71 +129,17 @@ float MobilenetDetection::iou(const tk::dnn::box &a, const tk::dnn::box &b)
     return iou;
 }
 
-std::vector<tk::dnn::box> MobilenetDetection::postprocess(const int width, const int height)
+bool MobilenetDetection::init(const std::string& tensor_path, const int n_classes)
 {
-    float *conf_per_class;
-    std::vector<tk::dnn::box> detections;
-    for (int i = 1; i < classes; i++){
-        conf_per_class = &confidences_h[i * nPriors];
-        std::vector<tk::dnn::box> boxes;
-        for (int j = 0; j < nPriors; j++){
 
-            if (conf_per_class[j] > confThreshold){
-                tk::dnn::box b;
-                b.cl = i;
-                b.prob = conf_per_class[j];
-                b.x = locations_h[j * N_COORDS + 0];
-                b.y = locations_h[j * N_COORDS + 1];
-                b.w = locations_h[j * N_COORDS + 2];
-                b.h = locations_h[j * N_COORDS + 3];
+    std::cout<<"MobilenetDetection Init"<<std::endl;
+    netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
+    imageSize = netRT->input_dim.h;
+    classes = n_classes;
 
-                boxes.push_back(b);
-            }
-        }
-        std::sort(boxes.begin(), boxes.end(), boxProbCmp);
+    SSDSpec specs[N_SSDSPEC];
 
-        std::vector<tk::dnn::box> remaining;
-        while (boxes.size() > 0){
-            remaining.clear();
-
-            tk::dnn::box b;
-            b.cl = boxes[0].cl;
-            b.prob = boxes[0].prob;
-            b.x = boxes[0].x * width;
-            b.y = boxes[0].y * height;
-            b.w = boxes[0].w * width;
-            b.h = boxes[0].h * height;
-            detections.push_back(b);
-            for (size_t j = 1; j < boxes.size(); j++){
-                if (iou(boxes[0], boxes[j]) <= IoUThreshold){
-                    remaining.push_back(boxes[j]);
-                }
-            }
-            boxes = remaining;
-        }
-    }
-    return detections;
-}
-
-float MobilenetDetection::get_color2(int c, int x, int max)
-{
-    float ratio = ((float)x / max) * 5;
-    int i = floor(ratio);
-    int j = ceil(ratio);
-    ratio -= i;
-    float r = (1 - ratio) * __colors[i % 6][c % 3] + ratio * __colors[j % 6][c % 3];
-    return r;
-}
-
-void MobilenetDetection::init(std::string tensor_path, int input_size, int n_classes)
-{
-    this->imageSize = input_size;
-    this->classes = n_classes;
-
-    const int n_SSDSpec = 6;
-    SSDSpec specs[6];
-
-    if(input_size == 300){
+    if(imageSize == 300){
         specs[0].setAll(19, 16, 60, 105, 2, 3);
         specs[1].setAll(10, 32, 105, 150, 2, 3);
         specs[2].setAll(5, 64, 150, 195, 2, 3);
@@ -203,7 +147,7 @@ void MobilenetDetection::init(std::string tensor_path, int input_size, int n_cla
         specs[4].setAll(2, 150, 240, 285, 2, 3);
         specs[5].setAll(1, 300, 285, 330, 2, 3);
     }
-    else if(input_size == 512){
+    else if(imageSize == 512){
         specs[0].setAll(32, 16, 60, 105, 2, 3);
         specs[1].setAll(16, 32, 105, 150, 2, 3);
         specs[2].setAll(8, 64, 150, 195, 2, 3);
@@ -215,23 +159,21 @@ void MobilenetDetection::init(std::string tensor_path, int input_size, int n_cla
         FatalError("Input size for mobilenet not supported");
     }
 
-    generate_ssd_priors(specs, n_SSDSpec);
-
-    netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
+    generate_ssd_priors(specs, N_SSDSPEC);
 
+#ifndef OPENCV_CUDA
     checkCuda(cudaMallocHost(&input, sizeof(dnnType) * netRT->input_dim.tot()));
+#endif
     checkCuda(cudaMalloc(&input_d, sizeof(dnnType) * netRT->input_dim.tot()));
 
     locations_h = (float *)malloc(N_COORDS * nPriors * sizeof(float));
     confidences_h = (float *)malloc(nPriors * classes * sizeof(float));
 
-    dim = tk::dnn::dataDim_t(1, 3, imageSize, imageSize, 1);
-
     for (int c = 0; c < classes; c++){
         int offset = c * 123457 % classes;
-        float r = get_color2(2, offset, classes);
-        float g = get_color2(1, offset, classes);
-        float b = get_color2(0, offset, classes);
+        float r = getColor(2, offset, classes);
+        float g = getColor(1, offset, classes);
+        float b = getColor(0, offset, classes);
         colors[c] = cv::Scalar(int(255.0 * b), int(255.0 * g), int(255.0 * r));
     }
 
@@ -263,99 +205,151 @@ void MobilenetDetection::init(std::string tensor_path, int input_size, int n_cla
     else{
         FatalError("Number of classes not supported for mobilenet");
     }
+    return 1;
 }
 
-cv::Mat MobilenetDetection::draw()
-{
-    tk::dnn::box b;
-    for (size_t i = 0; i < detected.size(); i++){
-        b = detected[i];
-        std::string det_class = classesNames[b.cl];
-        cv::rectangle(origImg, cv::Point(b.x, b.y), cv::Point(b.w, b.h), colors[b.cl], 2);
-        // draw label
-        cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
-        cv::rectangle(origImg, cv::Point(b.x, b.y), cv::Point((b.x + textSize.width - 2), (b.y - textSize.height - 2)), colors[b.cl], -1);
-        cv::putText(origImg, det_class, cv::Point(b.x, (b.y - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
-    }
-    return origImg;
-}
-
-void MobilenetDetection::preprocess()
+void MobilenetDetection::preprocess(cv::Mat &frame)
 {
     std::cout<<"preprocess"<<std::endl;
 #ifdef OPENCV_CUDA
         //move original image on GPU
-        cv::cuda::GpuMat im_Orig, frame_resize, frame_nomean, frame_scaled;
-        im_Orig = cv::cuda::GpuMat(origImg);
+        cv::cuda::GpuMat orig_img, frame_nomean;
+        orig_img = cv::cuda::GpuMat(frame);
 
         //resize image, remove mean, divide by std
-        cv::cuda::resize (im_Orig, frame_resize, cv::Size(netRT->input_dim.w, netRT->input_dim.h)); 
-        frame_resize.convertTo(frame_nomean, CV_32FC3, 1, -127);
-        frame_nomean.convertTo(frame_scaled, CV_32FC3, 1 / 128.0, 0);
+        cv::cuda::resize (orig_img, orig_img, cv::Size(netRT->input_dim.w, netRT->input_dim.h)); 
+        orig_img.convertTo(frame_nomean, CV_32FC3, 1, -127);
+        frame_nomean.convertTo(imagePreproc, CV_32FC3, 1 / 128.0, 0);
 
         //copy image into tensors
-        cv::cuda::GpuMat bgr[3];
-        cv::cuda::split(frame_scaled, bgr);
+        cv::cuda::split(imagePreproc, bgr);
 
         for(int i=0; i < netRT->input_dim.c; i++){
-            int idx = i * frame_scaled.rows * frame_scaled.cols;
-            checkCuda( cudaMemcpy((void *)&input_d[idx], (void *)bgr[i].data, frame_scaled.rows * frame_scaled.cols* sizeof(float), cudaMemcpyDeviceToDevice) );
+            int idx = i * imagePreproc.rows * imagePreproc.cols;
+            checkCuda( cudaMemcpy((void *)&input_d[idx], (void *)bgr[i].data, imagePreproc.rows * imagePreproc.cols* sizeof(float), cudaMemcpyDeviceToDevice) );
         }
 #else
         //resize image, remove mean, divide by std
-        cv::Mat frame_resize, frame_nomean, frame_scaled;
-        resize(origImg, frame_resize, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
-        frame_resize.convertTo(frame_nomean, CV_32FC3, 1, -127);
-        frame_nomean.convertTo(frame_scaled, CV_32FC3, 1 / 128.0, 0);
+        cv::Mat frame_nomean;
+        resize(frame, frame, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
+        frame.convertTo(frame_nomean, CV_32FC3, 1, -127);
+        frame_nomean.convertTo(imagePreproc, CV_32FC3, 1 / 128.0, 0);
 
         //copy image into tensor and copy it into GPU
-        cv::Mat bgr[3];
-        cv::split(frame_scaled, bgr);
+        cv::split(imagePreproc, bgr);
         for (int i = 0; i < netRT->input_dim.c; i++){
-            int idx = i * frame_scaled.rows * frame_scaled.cols;
-            memcpy((void *)&input[idx], (void *)bgr[i].data, frame_scaled.rows * frame_scaled.cols * sizeof(dnnType));
+            int idx = i * imagePreproc.rows * imagePreproc.cols;
+            memcpy((void *)&input[idx], (void *)bgr[i].data, imagePreproc.rows * imagePreproc.cols * sizeof(dnnType));
         }
         checkCuda(cudaMemcpyAsync(input_d, input, netRT->input_dim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
 #endif
 }
 
-void MobilenetDetection::update(cv::Mat &img)
+void MobilenetDetection::update(cv::Mat &frame)
 {
     TIMER_START
     detected.clear();
 
-    //save origin image
-    origImg = img;
-    cv::Size sz = origImg.size();
+    originalSize = frame.size();
 
     //preprocess
-    preprocess();
+    preprocess(frame);
 
     //do inference
-    tk::dnn::dataDim_t dim2 = dim;
+    tk::dnn::dataDim_t dim = tk::dnn::dataDim_t(1, 3, imageSize, imageSize, 1);;
     printCenteredTitle(" TENSORRT inference ", '=', 30);
     {
-        dim2.print();
+        dim.print();
         TIMER_START
-        netRT->infer(dim2, input_d);
+        netRT->infer(dim, input_d);
         TIMER_STOP
-        dim2.print();
+        dim.print();
     }
 
     //get confidences and locations_h
-    conf = (dnnType *)netRT->buffersRT[3];
-    loc = (dnnType *)netRT->buffersRT[4];
-
-    checkCuda(cudaMemcpy(confidences_h, conf, nPriors * classes * sizeof(float), cudaMemcpyDeviceToHost));
-    checkCuda(cudaMemcpy(locations_h, loc, N_COORDS * nPriors * sizeof(float), cudaMemcpyDeviceToHost));
+    dnnType *rt_out[2];
+    rt_out[0] = (dnnType *)netRT->buffersRT[3];
+    rt_out[1] = (dnnType *)netRT->buffersRT[4];
 
     //postprocess
-    convert_locatios_to_boxes_and_center();
-    detected = postprocess(sz.width, sz.height);
+    postprocess(rt_out, 2);
 
     TIMER_STOP
     stats.push_back(t_ns);
 }
 
+
+void MobilenetDetection::postprocess(dnnType **rt_out, const int n_out)
+{
+    checkCuda(cudaMemcpy(confidences_h, rt_out[0], nPriors * classes * sizeof(float), cudaMemcpyDeviceToHost));
+    checkCuda(cudaMemcpy(locations_h, rt_out[1], N_COORDS * nPriors * sizeof(float), cudaMemcpyDeviceToHost));
+    convert_locatios_to_boxes_and_center();
+
+    int width =  originalSize.width;
+    int height =  originalSize.height;
+
+    float *conf_per_class;
+    for (int i = 1; i < classes; i++){
+        conf_per_class = &confidences_h[i * nPriors];
+        std::vector<tk::dnn::box> boxes;
+        for (int j = 0; j < nPriors; j++){
+
+            if (conf_per_class[j] > confThreshold){
+                tk::dnn::box b;
+                b.cl = i;
+                b.prob = conf_per_class[j];
+                b.x = locations_h[j * N_COORDS + 0];
+                b.y = locations_h[j * N_COORDS + 1];
+                b.w = locations_h[j * N_COORDS + 2];
+                b.h = locations_h[j * N_COORDS + 3];
+
+                boxes.push_back(b);
+            }
+        }
+        std::sort(boxes.begin(), boxes.end(), boxProbCmp);
+
+        std::vector<tk::dnn::box> remaining;        
+        while (boxes.size() > 0){
+            remaining.clear();
+
+            tk::dnn::box b;
+            b.cl = boxes[0].cl;
+            b.prob = boxes[0].prob;
+            b.x = boxes[0].x * width;
+            b.y = boxes[0].y * height;
+            b.w = boxes[0].w * width;
+            b.h = boxes[0].h * height;
+            detected.push_back(b);
+            for (size_t j = 1; j < boxes.size(); j++){
+                if (iou(boxes[0], boxes[j]) <= IoUThreshold){
+                    remaining.push_back(boxes[j]);
+                }
+            }
+            boxes = remaining;
+        }
+    }
+}
+
+
+cv::Mat MobilenetDetection::draw(cv::Mat &frame)
+{
+    int baseline = 0;
+    float font_scale = 0.5;
+    int thickness = 2;
+    
+    tk::dnn::box b;
+    for (size_t i = 0; i < detected.size(); i++){
+        b = detected[i];
+        std::string det_class = classesNames[b.cl];
+        cv::rectangle(frame, cv::Point(b.x, b.y), cv::Point(b.w, b.h), colors[b.cl], 2);
+        // draw label
+        cv::Size text_size = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
+        cv::rectangle(frame, cv::Point(b.x, b.y), cv::Point((b.x + text_size.width - 2), (b.y - text_size.height - 2)), colors[b.cl], -1);
+        cv::putText(frame, det_class, cv::Point(b.x, (b.y - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
+    }
+    return frame;
+}
+
+
 } // namespace dnn
 } // namespace tk
\ No newline at end of file
diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp
index 08742c6..f94dd94 100644
--- a/src/Yolo3Detection.cpp
+++ b/src/Yolo3Detection.cpp
@@ -1,28 +1,14 @@
 #include "Yolo3Detection.h"
 
+
 namespace tk { namespace dnn {
 
-float _colors[6][3] = { {1,0,1}, {0,0,1},{0,1,1},{0,1,0},{1,1,0},{1,0,0} };
-float get_color(int c, int x, int max)
-{
-    float ratio = ((float)x/max)*5;
-    int i = floor(ratio);
-    int j = ceil(ratio);
-    ratio -= i;
-    float r = (1-ratio) * _colors[i % 6][c % 3] + ratio*_colors[j % 6][c % 3];
-    //printf("%f\n", r);
-    return r;
-}
-
-bool Yolo3Detection::init(std::string tensor_path) {
-
-    //const char *tensor_path = "../data/yolo3/yolo3_berkeley.rt";
+bool Yolo3Detection::init(const std::string& tensor_path, const int n_classes) {
 
     //convert network to tensorRT
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
 
-
     if(netRT->pluginFactory->n_yolos < 2 ) {
         FatalError("this is not yolo3");
     }
@@ -31,119 +17,124 @@ bool Yolo3Detection::init(std::string tensor_path) {
         YoloRT *yRT = netRT->pluginFactory->yolos[i];
         classes = yRT->classes;
         num = yRT->num;
-        n_masks = yRT->n_masks;
+        nMasks = yRT->n_masks;
 
         // make a yolo layer for interpret predictions
-        yolo[i] = new tk::dnn::Yolo(nullptr, classes, n_masks, ""); // yolo without input and bias
-        yolo[i]->mask_h = new dnnType[n_masks];
-        yolo[i]->bias_h = new dnnType[num*n_masks*2];
-        memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*n_masks);
-        memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*n_masks*2);
+        yolo[i] = new tk::dnn::Yolo(nullptr, classes, nMasks, ""); // yolo without input and bias
+        yolo[i]->mask_h = new dnnType[nMasks];
+        yolo[i]->bias_h = new dnnType[num*nMasks*2];
+        memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*nMasks);
+        memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*nMasks*2);
         yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w);
         yolo[i]->classesNames = yRT->classesNames;
     }
 
     dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
-    
+#ifndef OPENCV_CUDA
     checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot()));
+#endif
     checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot()));
 
-
     // class colors precompute    
     for(int c=0; c<classes; c++) {
         int offset = c*123457 % classes;
-        float r = get_color(2, offset, classes);
-        float g = get_color(1, offset, classes);
-        float b = get_color(0, offset, classes);
+        float r = getColor(2, offset, classes);
+        float g = getColor(1, offset, classes);
+        float b = getColor(0, offset, classes);
         colors[c] = cv::Scalar(int(255.0*b), int(255.0*g), int(255.0*r));
     }
     return true;
 } 
 
-cv::Mat Yolo3Detection::draw(cv::Mat &imageORIG) {
+void Yolo3Detection::preprocess(cv::Mat &frame)
+{
+#ifdef OPENCV_CUDA
+    cv::cuda::GpuMat orig_img, img_resized;
+    orig_img = cv::cuda::GpuMat(frame);
+    cv::cuda::resize(orig_img, img_resized, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
 
-    tk::dnn::box b;
-    int x0, w, x1, y0, h, y1;
-    int objClass;
-    std::string det_class;
-    float prob;
-    int baseline = 0;
-    float fontScale = 0.5;
-    int thickness = 2;   
-    // draw dets
-    for(int i=0; i<detected.size(); i++) {
-        b           = detected[i];
-        x0   		= b.x;
-        x1   		= b.x + b.w;
-        y0   		= b.y;
-        y1   		= b.y + b.h;
-        det_class 	= getYoloLayer()->classesNames[b.cl];
-        prob    	= b.prob;
-
-        // std::cout<<det_class<<" ("<<prob<<"): "<<x0<<" "<<y0<<" "<<x1<<" "<<y1<<"\n";
-        // draw rectangle
-        cv::rectangle(imageORIG, cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2); 
-
-        // draw label
-        cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
-        cv::rectangle(imageORIG, cv::Point(x0, y0), cv::Point((x0 + textSize.width - 2), (y0 - textSize.height - 2)), colors[b.cl], -1);                      
-        cv::putText(imageORIG, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
-    }
-    return imageORIG;
-}
-
-void Yolo3Detection::update(cv::Mat &imageORIG) {
-
-    TIMER_START
-    if(!imageORIG.data) {
-        std::cout<<"YOLO: NO IMAGE DATA\n";
-        return;
-    }     
-    float xRatio =  float(imageORIG.cols) / float(netRT->input_dim.w);
-    float yRatio =  float(imageORIG.rows) / float(netRT->input_dim.h);
-
-    resize(imageORIG, imageORIG, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
-
-    imageORIG.convertTo(imageF, CV_32FC3, 1/255.0); 
+    img_resized.convertTo(imagePreproc, CV_32FC3, 1/255.0); 
 
     //split channels
-    cv::split(imageF,bgr);//split source
+    cv::cuda::split(imagePreproc,bgr);//split source
 
     //write channels
     for(int i=0; i<netRT->input_dim.c; i++) {
-        int idx = i*imageF.rows*imageF.cols;
+        std::cout<<"copio il channel"<<i<<std::endl;
+        int idx = i*imagePreproc.rows*imagePreproc.cols;
         int ch = netRT->input_dim.c-1 -i;
-        memcpy((void*)&input[idx], (void*)bgr[ch].data, imageF.rows*imageF.cols*sizeof(dnnType));
+        checkCuda( cudaMemcpy((void*)&input_d[idx], (void*)bgr[ch].data, imagePreproc.rows*imagePreproc.cols*sizeof(dnnType), cudaMemcpyDeviceToDevice));
     }
+#else
+    cv::resize(frame, frame, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
+    frame.convertTo(imagePreproc, CV_32FC3, 1/255.0); 
 
+    //split channels
+    cv::split(imagePreproc,bgr);//split source
 
-    //DO INFERENCE
-    dnnType *rt_out[netRT->pluginFactory->n_yolos]; 
+    //write channels
+    for(int i=0; i<netRT->input_dim.c; i++) {
+        int idx = i*imagePreproc.rows*imagePreproc.cols;
+        int ch = netRT->input_dim.c-1 -i;
+        memcpy((void*)&input[idx], (void*)bgr[ch].data, imagePreproc.rows*imagePreproc.cols*sizeof(dnnType));
+    }
+    checkCuda(cudaMemcpyAsync(input_d, input, netRT->input_dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
+#endif
+}
+
+void Yolo3Detection::update(cv::Mat &frame)
+{
+    TIMER_START
+    if(!frame.data) {
+        std::cout<<"YOLO: NO IMAGE DATA\n";
+        return;
+    }   
+
+    originalSize = frame.size();
+    preprocess(frame);
+
+    //do inference
     tk::dnn::dataDim_t dim = netRT->input_dim;
-    checkCuda(cudaMemcpyAsync(input_d, input, dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
-
-    printCenteredTitle(" TENSORRT inference ", '=', 30); {
+    
+    printCenteredTitle(" TENSORRT inference ", '=', 30); 
+    {
         dim.print();
         TIMER_START
         netRT->infer(dim, input_d);
         TIMER_STOP
         dim.print();
-    
-        stats.push_back(t_ns);
     }
-    
-    // compute dets
-    ndets = 0;
+
+    //get yolo outputs
+    dnnType *rt_out[netRT->pluginFactory->n_yolos]; 
     for(int i=0; i<netRT->pluginFactory->n_yolos; i++) {
         rt_out[i] = (dnnType*)netRT->buffersRT[i+1];
-        yolo[i]->dstData = rt_out[i];
-        yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, thresh);
     }
-    tk::dnn::Yolo::mergeDetections(dets, ndets, classes);
+
+    postprocess(rt_out, netRT->pluginFactory->n_yolos);
+        
+    TIMER_STOP
+    stats.push_back(t_ns);
+}
+
+void Yolo3Detection::postprocess(dnnType **rt_out, const int n_out)
+{
+    float x_ratio =  float(originalSize.width) / float(netRT->input_dim.w);
+    float y_ratio =  float(originalSize.height) / float(netRT->input_dim.h);
+
+    std::cout<<"RATIO:"<<x_ratio<<" "<<y_ratio<<std::endl;
+
+    // compute dets
+    nDets = 0;
+    for(int i=0; i<n_out; i++) {
+        yolo[i]->dstData = rt_out[i];
+        yolo[i]->computeDetections(dets, nDets, netRT->input_dim.w, netRT->input_dim.h, confThreshold);
+    }
+    tk::dnn::Yolo::mergeDetections(dets, nDets, classes);
 
     // fill detected
     detected.clear();
-    for(int j=0; j<ndets; j++) {
+    for(int j=0; j<nDets; j++) {
         tk::dnn::Yolo::box b = dets[j].bbox;
         int x0   = (b.x-b.w/2.);
         int x1   = (b.x+b.w/2.);
@@ -152,21 +143,18 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
         int obj_class = -1;
         float prob = 0;
         for(int c=0; c<classes; c++) {
-            if(dets[j].prob[c] >= thresh) {
+            if(dets[j].prob[c] >= confThreshold) {
                 obj_class = c;
                 prob = dets[j].prob[c];
             }
         }
 
         if(obj_class >= 0) {
-            //std::cout<<obj_class<<" ("<<prob<<"): "<<x0<<" "<<y0<<" "<<x1<<" "<<y1<<"\n";
-            //cv::rectangle(image, cv::Point(x0, y0), cv::Point(x1, y1), colors[obj_class], 2);
-            
             // convert to image coords
-            x0 = xRatio*x0;
-            x1 = xRatio*x1;
-            y0 = yRatio*y0;
-            y1 = yRatio*y1;
+            x0 = x_ratio*x0;
+            x1 = x_ratio*x1;
+            y0 = y_ratio*y0;
+            y1 = y_ratio*y1;
               
             tk::dnn::box res;
             res.cl = obj_class;
@@ -178,9 +166,44 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
             detected.push_back(res);
         }
     }
-    TIMER_STOP
-    stats.push_back(t_ns);
+    std::cout<<"N detections: "<<detected.size()<<std::endl;
+}
 
+cv::Mat Yolo3Detection::draw(cv::Mat &frame) 
+{
+    tk::dnn::box b;
+    int x0, w, x1, y0, h, y1;
+    int objClass;
+    std::string det_class;
+    int baseline = 0;
+    float font_scale = 0.5;
+    int thickness = 2;   
+    // draw dets
+    for(int i=0; i<detected.size(); i++) {
+        b           = detected[i];
+        x0   		= b.x;
+        x1   		= b.x + b.w;
+        y0   		= b.y;
+        y1   		= b.y + b.h;
+        det_class 	= getYoloLayer()->classesNames[b.cl];
+
+        // draw rectangle
+        cv::rectangle(frame, cv::Point(x0, y0), cv::Point(x1, y1), colors[b.cl], 2); 
+
+        // draw label
+        cv::Size text_size = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
+        cv::rectangle(frame, cv::Point(x0, y0), cv::Point((x0 + text_size.width - 2), (y0 - text_size.height - 2)), colors[b.cl], -1);                      
+        cv::putText(frame, det_class, cv::Point(x0, (y0 - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
+    }
+    return frame;
+}
+
+tk::dnn::Yolo* Yolo3Detection::getYoloLayer(int n) 
+{
+    if(n<3)
+        return yolo[n];
+    else 
+        return nullptr;
 }
 
 }}
diff --git a/src/utils.cpp b/src/utils.cpp
index 142e85f..2606741 100644
--- a/src/utils.cpp
+++ b/src/utils.cpp
@@ -129,6 +129,18 @@ int checkResult(int size, dnnType *data_d, dnnType *correct_d, bool device) {
     return diffs;
 }
 
+float getColor(const int c, const int x, const int max)
+{
+    float _colors[6][3] = { {1,0,1}, {0,0,1},{0,1,1},{0,1,0},{1,1,0},{1,0,0} };
+    float ratio = ((float)x/max)*5;
+    int i = floor(ratio);
+    int j = ceil(ratio);
+    ratio -= i;
+    float r = (1-ratio) * _colors[i % 6][c % 3] + ratio*_colors[j % 6][c % 3];
+    return r;
+}
+
+
 void resize(int size, dnnType **data)
 {
     if (*data != NULL)