YoloRT save bias, mask and clasesName into RT file

include/tkDNN/CenternetDetection.h

@@ -73,7 +73,7 @@ public:
     CenternetDetection() {};
     ~CenternetDetection() {}; 
 
-    bool init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
+    bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
     void preprocess(cv::Mat &frame, const int bi=0);
     void postprocess(const int bi=0,const bool mAP=false);
 };

include/tkDNN/DetectionNN.h

@@ -87,7 +87,7 @@ class DetectionNN {
          * @param n_batches maximum number of batches to use in inference
          * @return true if everything is correct, false otherwise.
          */
-        virtual bool init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3) = 0;
+        virtual bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3) = 0;
         
         /**
          * This method performs the whole detection of the NN.

include/tkDNN/MobilenetDetection.h

@@ -65,7 +65,7 @@ public:
     MobilenetDetection() {};
     ~MobilenetDetection() {}; 
 
-    bool init(const std::string& tensor_path, const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches=1, const float conf_thresh=0.3);
+    bool init(const std::string& tensor_path,const int n_classes, const int n_batches=1, const float conf_thresh=0.3);
     void preprocess(cv::Mat &frame, const int bi=0);
     void postprocess(const int bi=0,const bool mAP=false);
 };

include/tkDNN/NetworkRT.h

@@ -30,10 +30,6 @@
 
 namespace tk { namespace dnn {
 
-using namespace nvinfer1;
-
-
-
 class NetworkRT {
 
 public:
@@ -57,6 +53,7 @@ public:
     dnnType *output;
     cudaStream_t stream;
 
+    std::vector<nvinfer1::YoloRT*> yolo_plugins; // yolo layers in network
 
     NetworkRT(Network *net, const char *name);
     virtual ~NetworkRT();

include/tkDNN/Yolo3Detection.h

@@ -19,13 +19,12 @@ private:
     tk::dnn::Yolo* getYoloLayer(int n=0);
 
     cv::Mat bgr_h;
-    std::vector<int> noYolos;
     
 public:
     Yolo3Detection() {};
     ~Yolo3Detection() {}; 
 
-    bool init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path,const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
+    bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
     void preprocess(cv::Mat &frame, const int bi=0);
     void postprocess(const int bi=0,const bool mAP=false);
 };

include/tkDNN/pluginsRT/YoloRT.h

@@ -5,7 +5,6 @@
 #include <vector>
 #include "../kernels.h"
 #include <NvInfer.h>
-#include <tkdnn.h>
 
 #define YOLORT_CLASSNAME_W 256
 
@@ -80,8 +79,10 @@ namespace nvinfer1 {
         float nms_thresh;
         int nms_kind;
         int new_coords;
-        int NUM = 0;
+
         std::vector<std::string> classesNames;
+        std::vector<dnnType> mask;
+        std::vector<dnnType> bias;
 
 
         int entry_index(int batch, int location, int entry) {