yolo3plug fix

src/NetworkRT.cpp

@@ -360,7 +360,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Yolo *l) {
     //std::cout<<"convert Yolo\n";
 
     //std::cout<<"New plugin YOLO\n";
-    IPlugin *plugin = new YoloRT(l->classes, l->num);
+    IPlugin *plugin = new YoloRT(l->classes, l->num, l);
     IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
     checkNULL(lRT);
     return lRT;
@@ -396,6 +396,7 @@ bool NetworkRT::serialize(const char *filename) {
 class PluginFactory : IPluginFactory
 {
 public:
+
 	virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) {
         const char * buf = reinterpret_cast<const char*>(serialData);
         
@@ -440,6 +441,12 @@ public:
         	r->c = readBUF<int>(buf);
 		    r->h = readBUF<int>(buf);
 		    r->w = readBUF<int>(buf);
+            for(int i=0; i<r->num; i++)
+        		r->mask[i] = readBUF<dnnType>(buf);
+            for(int i=0; i<3*2*r->num; i++)
+        		r->bias[i] = readBUF<dnnType>(buf);
+
+            std::cout<<"YOLO: "<<r->c<<" "<<r->h<<" "<<r->w<<"\n";
             return r;
         } 
 

src/Yolo3Detection.cpp

@@ -21,14 +21,14 @@ bool Yolo3Detection::init(std::string tensor_folder) {
     }
 
     //convert network to tensorRT
-    std::cout<<(tensor_folder + "/yolo3_berkeley.rt").c_str()<<"\n";
-    netRT = new tk::dnn::NetworkRT(NULL, (tensor_folder + "/yolo3_berkeley.rt").c_str() );
+    std::cout<<(tensor_folder + ".rt").c_str()<<"\n";
+    netRT = new tk::dnn::NetworkRT(NULL, (tensor_folder + ".rt").c_str() );
 
-    yolo[0] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "/yolo3_0.bin").c_str() ); // yolo without input and bias
+    yolo[0] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "_0.bin").c_str() ); // yolo without input and bias
     yolo[0]->input_dim = yolo[0]->output_dim = tk::dnn::dataDim_t(1, 45, 10, 17);
-    yolo[1] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "/yolo3_1.bin").c_str() ); // yolo without input and bias
+    yolo[1] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "_1.bin").c_str() ); // yolo without input and bias
     yolo[1]->input_dim = yolo[1]->output_dim = tk::dnn::dataDim_t(1, 45, 20, 34);
-    yolo[2] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "/yolo3_2.bin").c_str() ); // yolo without input and bias
+    yolo[2] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "_2.bin").c_str() ); // yolo without input and bias
     yolo[2]->input_dim = yolo[2]->output_dim = tk::dnn::dataDim_t(1, 45, 40, 68);
 
     dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);

src/pluginsRT/YoloRT.cpp

@@ -3,11 +3,20 @@
 
 class YoloRT : public IPlugin {
 
+
+
 public:
-	YoloRT(int classes, int num) {
+	YoloRT(int classes, int num, tk::dnn::Yolo *yolo = nullptr) {
 
 		this->classes = classes;
 		this->num = num;
+
+        mask = new dnnType[num];
+        bias = new dnnType[num*3*2];
+        if(yolo != nullptr) {
+            memcpy(mask, yolo->mask_h, sizeof(dnnType)*num);
+            memcpy(bias, yolo->bias_h, sizeof(dnnType)*num*3*2);
+        }
 	}
 
 	~YoloRT(){
@@ -63,7 +72,7 @@ public:
 
 
 	virtual size_t getSerializationSize() override {
-		return 5*sizeof(int);
+		return 5*sizeof(int) + num*sizeof(dnnType) + num*3*2*sizeof(dnnType);
 	}
 
 	virtual void serialize(void* buffer) override {
@@ -73,11 +82,18 @@ public:
 		tk::dnn::writeBUF(buf, c);
 		tk::dnn::writeBUF(buf, h);
 		tk::dnn::writeBUF(buf, w);
+        for(int i=0; i<num; i++)
+    		tk::dnn::writeBUF(buf, mask[i]);
+        for(int i=0; i<3*2*num; i++)
+    		tk::dnn::writeBUF(buf, bias[i]);
 	}
 
 	int c, h, w;
     int classes, num;
 
+    dnnType *mask;
+    dnnType *bias;
+
 	int entry_index(int batch, int location, int entry, int batchSize) {
 		int n =   location / (w*h);
 		int loc = location % (w*h);