YoloRT save bias, mask and clasesName into RT file

demo/demo/demo.cpp

@@ -39,12 +39,8 @@ int main(int argc, char *argv[]) {
 
     #ifdef __linux__
         std::string input = YAMLgetConf<std::string>(conf, "input", "../demo/yolo_test.mp4");
-        std::string cfgPath = YAMLgetConf<std::string>(conf,"cfg_input", "../tests/darknet/cfg/yolo4tiny.cfg");
-        std::string namePath = YAMLgetConf<std::string>(conf,"name_input","../tests/darknet/names/coco.names");
     #elif _WIN32
         std::string input = YAMLgetConf<std::string>(conf, "win_input", "..\\..\\..\\demo\\yolo_test.mp4");
-        std::string cfgPath = YAMLgetConf<std::string>(conf,"cfg_win_input","..\\..\\..\\tests\\darknet\\cfg\\yolo4tiny.cfg");
-        std::string namePath = YAMLgetConf<std::string>(conf,"name_win_input","..\\..\\..\\tests\\darknet\\names\\coco.names");
     #endif
         if(!fileExist(input.c_str()))
         FatalError("The given input video does not exist.");
@@ -90,12 +86,7 @@ int main(int argc, char *argv[]) {
         FatalError("Network type not allowed (3rd parameter)\n");
     }
 
-    if(ntype == 'c' || ntype == 'm'){
-        cfgPath = "";
-        namePath = "";
-    }
-
-    detNN->init(net,cfgPath,namePath,n_classes,n_batch,conf_thresh);
+    detNN->init(net,n_classes,n_batch,conf_thresh);
 
     // open video stream
     cv::VideoCapture cap(input);

demo/demo/map.cpp

@@ -45,8 +45,6 @@ int main(int argc, char *argv[])
     bool verbose;
     int classes, map_points, map_levels;
     float map_step, IoU_thresh, conf_thresh;
-    std::string cfg_path = "../tests/darknet/cfg/yolo4tiny.cfg";
-    std::string name_path = "../tests/darknet/names/coco.names";
     double vm_total = 0, rss_total = 0;
     double vm, rss;
 
@@ -56,17 +54,13 @@ int main(int argc, char *argv[])
     if(argc > 2)
         ntype = argv[2][0];
     if(argc > 3)
-        cfg_path = argv[3];
+        labels_path = argv[3]; 
     if(argc > 4)
-        name_path = argv[4];
+        config_filename = argv[4]; 
     if(argc > 5)
-        labels_path = argv[5];
+        n_batches = atoi(argv[5]); 
     if(argc > 6)
-        config_filename = argv[6];
-    if(argc > 7)
-        n_batches = atoi(argv[7]);
-    if(argc > 8)
-        confidence_thresh = atof(argv[8]);
+        confidence_thresh = atof(argv[6]); 
 
     std::cout<<"conf t: "<<confidence_thresh<<std::endl;
 
@@ -121,7 +115,7 @@ int main(int argc, char *argv[])
         default:
             FatalError("Network type not allowed (3rd parameter)\n");
     }
-    detNN->init(net,cfg_path,name_path,n_classes, 1, conf_thresh);
+    detNN->init(net,n_classes, 1, conf_thresh);
 
     //read images 
     std::ifstream all_labels(labels_path);

demo/demoConfig.yaml

@@ -2,16 +2,8 @@
 input     : "../demo/yolo_test.mp4"
 win_input : "..\\..\\..\\demo\\yolo_test.mp4"
 
-#cfg input
-cfg_input : "../tests/darknet/cfg/yolo4tiny.cfg"
-cfg_win_input : "..\\..\\..\\tests\\darknet\\cfg\\yolo4tiny.cfg"
-
-#name input
-name_input : "../tests/darknet/names/coco.names"
-name_win_input : "..\\..\\..\\tests\\darknet\\names\\coco.names"
-
 # network config
-net         : "yolo4tiny_fp32.rt"
+net         : "yolo4_berkeley_fp32.rt"
 ntype       : 'y'
 n_classes   : 80
 n_batch     : 1

docs/demo.md

@@ -46,8 +46,6 @@ The config file is a yaml file with the following attributes:
 * ```conf_thresh``` confidence threshold for the detector. Only bounding boxes with threshold greater than conf-thresh will be displayed.
 * ```show``` if set to 0 the demo will not show the visualization (if n-batches ==1)
 * ```save``` if set to 1 the demo will save the video of the demo into result.mp4 (if n-batches ==1)
-* ```cfg_input``` (for linux) \ ```cfg_win_input``` (for windows)  is the location of the cfg path of the network for mobilenet and centernet networks use ```" "```
-* ```name_input``` (for linux) \ ```name_win_input``` (for windows)  is the location of the name path of the network for mobilenet and centernet networks use ```" "```
 
 N.B. By default it is used FP32 inference
 

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) {

src/CenternetDetection.cpp

@@ -3,7 +3,7 @@
 
 namespace tk { namespace dnn {
 
-bool CenternetDetection::init(const std::string& tensor_path, const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches, const float conf_thresh){
+bool CenternetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh){
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
     classes = n_classes;

src/MobilenetDetection.cpp

@@ -126,7 +126,7 @@ float MobilenetDetection::iou(const tk::dnn::box &a, const tk::dnn::box &b){
     return iou;
 }
 
-bool MobilenetDetection::init(const std::string& tensor_path, const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches, const float conf_thresh){
+bool MobilenetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh){
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
     imageSize = netRT->input_dim.h;

src/NetworkRT.cpp

@@ -15,6 +15,9 @@
 
 using namespace nvinfer1;
 
+extern std::mutex gYoloPlugins_mutex;
+extern std::vector<YoloRT*> gYoloPlugins;
+
 // Logger for info/warning/errors
 class Logger : public ILogger {
     void log(Severity severity, const char* msg) NOEXCEPT override {
@@ -826,6 +829,12 @@ IPluginV2Layer* NetworkRT::convert_layer(ITensor *input, Yolo *l) {
     mPluginAttributes.emplace_back(PluginField("nms_thresh",&l->nms_thresh,PluginFieldType::kFLOAT32,1));
     mPluginAttributes.emplace_back(PluginField("nms_kins",&l->nsm_kind,PluginFieldType::kINT32,1));
     mPluginAttributes.emplace_back(PluginField("new_coords",&l->new_coords,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("mask",l->mask_h,PluginFieldType::kFLOAT32,l->n_masks));
+    mPluginAttributes.emplace_back(PluginField("bias",l->bias_h,PluginFieldType::kFLOAT32,l->n_masks*2*l->num));
+    for(int i=0; i<l->classes; i++) {
+        mPluginAttributes.emplace_back(PluginField("class_name",l->classesNames[i].data(),PluginFieldType::kCHAR,l->classesNames[i].size()));
+    }
+
     mFC.nbFields = mPluginAttributes.size();
     mFC.fields = mPluginAttributes.data();
     auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC);
@@ -1001,7 +1010,14 @@ bool NetworkRT::deserialize(const char *filename) {
     }
 
     runtimeRT = createInferRuntime(loggerRT);
+
+    gYoloPlugins_mutex.lock();
+    gYoloPlugins.clear();
     engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size);
+    yolo_plugins = gYoloPlugins;
+    gYoloPlugins.clear();
+    gYoloPlugins_mutex.unlock();
+
     std::cout<<size<<std::endl;
     //if (gieModelStream) delete [] gieModelStream;
 

src/Yolo3Detection.cpp

@@ -3,7 +3,7 @@
 
 namespace tk { namespace dnn {
 
-    bool Yolo3Detection::init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches, const float conf_thresh) {
+    bool Yolo3Detection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh) {
 
     //convert network to tensorRT
     std::cout<<(tensor_path).c_str()<<"\n";
@@ -14,43 +14,27 @@ namespace tk { namespace dnn {
     tk::dnn::dataDim_t idim = netRT->input_dim;    
     idim.n = nBatches;
 
-    std::vector<int> yolosLine = noYolosLine(cfg_path);
-    noYolos = yolosLine;
-    int channels,height,width;
-    loadYoloInitInfo(channels,width,height,cfg_path);
-
-
-
-    if(yolosLine.size() < 2 ) {
+    if(netRT->yolo_plugins.size() < 2 ) {
         FatalError("this is not yolo3");
     }
 
+    for(int i=0; i<netRT->yolo_plugins.size(); i++) {
+        nvinfer1::YoloRT *yRT = netRT->yolo_plugins[i];
+        classes = yRT->classes;
+        num = yRT->num;
+        nMasks = yRT->n_masks;
 
-    for(int i=0; i<noYolos.size(); i++) {
-        std::vector<float> maskTemp,anchorsTemp;
-        std::vector<std::string> classNamesTemp;
-        int nms_kind,coords,numTemp;
-        float nmsthresh;
-        loadYoloInfo(cfg_path,yolosLine[i],maskTemp,anchorsTemp,numTemp,classes,nmsthresh,nms_kind,coords);
-        classNamesTemp = darknetReadNames(name_path);
-        num = numTemp/maskTemp.size();
-        nMasks = maskTemp.size();
-        dnnType* maskTempF;
-        dnnType* biasTempF;
-        maskTempF = maskTemp.data();
-        biasTempF = anchorsTemp.data();
         // make a yolo layer to interpret predictions
         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, maskTempF, sizeof(dnnType)*nMasks);
-        memcpy(yolo[i]->bias_h, biasTempF, sizeof(dnnType)*num*nMasks*2);
-        auto dim = netRT->engineRT->getBindingDimensions(i+1);
-        yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, dim.d[0], dim.d[1], dim.d[2]);
-        yolo[i]->classesNames = classNamesTemp;
-        yolo[i]->nms_thresh = nmsthresh;
-        yolo[i]->nsm_kind = (tk::dnn::Yolo::nmsKind_t) nms_kind;
-        yolo[i]->new_coords = coords;
+        memcpy(yolo[i]->mask_h, yRT->mask.data(), sizeof(dnnType)*nMasks);
+        memcpy(yolo[i]->bias_h, yRT->bias.data(), 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;
+        yolo[i]->nms_thresh = yRT->nms_thresh;
+        yolo[i]->nsm_kind = (tk::dnn::Yolo::nmsKind_t) yRT->nms_kind;
+        yolo[i]->new_coords = yRT->new_coords;
     }
 
     dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
@@ -112,12 +96,12 @@ void Yolo3Detection::postprocess(const int bi, const bool mAP){
 
 
     //get yolo outputs
-    if(noYolos.size() < 2){
+    if(netRT->yolo_plugins.size() < 2){
         FatalError("YOLOS WRONG!!");
     }
     std::vector<float *> rt_out;
     //dnnType *rt_out[netRT->pluginFactory->n_yolos];
-    for(int i=0; i<noYolos.size(); i++)
+    for(int i=0; i<netRT->yolo_plugins.size(); i++)
         rt_out.push_back((dnnType*)netRT->buffersRT[i+1] + netRT->buffersDIM[i+1].tot()*bi);
 
     float x_ratio =  float(originalSize[bi].width) / float(netRT->input_dim.w);
@@ -125,7 +109,7 @@ void Yolo3Detection::postprocess(const int bi, const bool mAP){
 
     // compute dets
     nDets = 0;
-    for(int i=0; i<noYolos.size(); i++) {
+    for(int i=0; i<netRT->yolo_plugins.size(); i++) {
         yolo[i]->dstData = rt_out[i];
         yolo[i]->computeDetections(dets, nDets, netRT->input_dim.w, netRT->input_dim.h, confThreshold, yolo[i]->new_coords);
     }

src/pluginsRT/YoloRT.cpp

@@ -1,8 +1,13 @@
 #include <tkDNN/pluginsRT/YoloRT.h>
 
 #include <utility>
+#include <mutex>
 using namespace nvinfer1;
 
+// used to retrive Yolo plugin during network deserialization
+std::mutex gYoloPlugins_mutex;
+std::vector<YoloRT*> gYoloPlugins;
+
 std::vector<PluginField> YoloRTPluginCreator::mPluginAttributes;
 PluginFieldCollection YoloRTPluginCreator::mFC{};
 
@@ -22,6 +27,10 @@ YoloRT::YoloRT(int classes, int num, int c,int h,int w,int n_masks, float scale_
     this->nms_thresh = nms_thresh;
     this->nms_kind = nms_kind;
     this->new_coords = new_coords;
+
+    bias.clear();
+    mask.clear();
+    classesNames.clear();
 }
 
 YoloRT::YoloRT(const void *data, size_t length) {
@@ -36,7 +45,24 @@ YoloRT::YoloRT(const void *data, size_t length) {
     c = readBUF<int>(buf);
     h = readBUF<int>(buf);
     w = readBUF<int>(buf);
+
+    mask.resize(n_masks);
+    for(int i=0; i<n_masks; i++)
+        mask[i] = readBUF<dnnType>(buf);
+    bias.resize(n_masks*2*num);
+    for(int i=0; i<n_masks*2*num; i++)
+        bias[i] = readBUF<dnnType>(buf);
+
+    // save classes names
+    classesNames.resize(classes);
+    for(int i=0; i<classes; i++) {
+        char tmp[YOLORT_CLASSNAME_W];
+        for(int j=0; j<YOLORT_CLASSNAME_W; j++)
+            tmp[j] = readBUF<char>(buf);
+        classesNames[i] = std::string(tmp);
+    }
     assert(buf == bufCheck + length);
+    gYoloPlugins.push_back(this);
 }
 
 YoloRT::~YoloRT() {}
@@ -126,7 +152,7 @@ int32_t YoloRT::enqueue(int32_t batchSize, const void *const *inputs, void **out
 
 
 size_t YoloRT::getSerializationSize() const NOEXCEPT {
-    return 8 * sizeof(int) + 2 * sizeof(float) ;
+    return 8 * sizeof(int) + 2 * sizeof(float) + n_masks*sizeof(dnnType) + num*n_masks*2*sizeof(dnnType) + YOLORT_CLASSNAME_W*classes*sizeof(char);
 }
 
 bool YoloRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT {
@@ -145,6 +171,19 @@ void YoloRT::serialize(void *buffer) const NOEXCEPT {
     writeBUF(buf, c);            //std::cout << "C : " << c << std::endl;
     writeBUF(buf, h);            //std::cout << "H : " << h << std::endl;
     writeBUF(buf, w);            //std::cout << "C : " << c << std::endl;
+    for (int i = 0; i < n_masks; i++)
+        writeBUF(buf, mask[i]); //std::cout << "mask[i] : " << mask[i] << std::endl;
+    for (int i = 0; i < n_masks * 2 * num; i++)
+        writeBUF(buf, bias[i]); //std::cout << "bias[i] : " << bias[i] << std::endl;
+
+    // save classes names
+    for(int i=0; i<classes; i++) {
+        char tmp[YOLORT_CLASSNAME_W];
+        strcpy(tmp, classesNames[i].c_str());
+        for(int j=0; j<YOLORT_CLASSNAME_W; j++) {
+            writeBUF(buf, tmp[j]);
+        }
+    }
 
     assert(buf == a + getSerializationSize());
 }
@@ -171,6 +210,9 @@ void YoloRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
 
 IPluginV2Ext *YoloRT::clone() const NOEXCEPT {
     auto *p = new YoloRT(classes, num,c,h,w,n_masks, scaleXY, nms_thresh, nms_kind, new_coords);
+    p->mask = mask;
+    p->bias = bias;
+    p->classesNames = classesNames;
     p->setPluginNamespace(mPluginNamespace.c_str());
     return p;
 }
@@ -235,6 +277,17 @@ IPluginV2Ext *YoloRTPluginCreator::createPlugin(const char *name, const PluginFi
     int nms_kind = *(static_cast<const int*>(fields[8].data));
     int new_coords = *(static_cast<const int*>(fields[9].data));
     auto *pluginObj = new YoloRT(classes,num,c,h,w,n_masks,scaleXY,nmsThresh,nms_kind,new_coords);
+
+    // fill additional data
+    pluginObj->mask.resize(fields[10].length*sizeof(float));
+    memcpy(pluginObj->mask.data(), fields[10].data, fields[10].length*sizeof(float));
+    pluginObj->bias.resize(fields[11].length*sizeof(float));
+    memcpy(pluginObj->bias.data(), fields[11].data, fields[11].length*sizeof(float));
+    pluginObj->classesNames.resize(classes);
+    for(int i=0; i<classes; i++) {
+        pluginObj->classesNames[i].resize(fields[12+i].length);
+        memcpy(&pluginObj->classesNames[i][0], fields[12+i].data, fields[12+i].length*sizeof(char));
+    }
     return pluginObj;
 }