save layer names in rt file

demo/demo/demo.cpp

@@ -12,10 +12,6 @@
 #include "Yolo3Detection.h"
 
 bool 			gRun;
-//std::string   	obj_class[10] {"person", "car", "truck", "bus", "motor", "bike", "rider", "traffic light", "traffic sign", "train"};
-//std::string   	obj_class[3] {"person", "bike", "car"};
-std::string   	obj_class[10] {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9"};
-
 bool SAVE_RESULT = false;
 
 void sig_handler(int signo) {
@@ -77,7 +73,7 @@ int main(int argc, char *argv[]) {
             int x1   				= b.x + b.w;
             int y0   				= b.y;
             int y1   				= b.y + b.h;
-            std::string det_class 	= obj_class[b.cl];
+            std::string det_class 	= yolo.getYoloLayer()->classesNames[b.cl];
             float prob 				= b.prob;
 
             std::cout<<det_class<<" ("<<prob<<"): "<<x0<<" "<<y0<<" "<<x1<<" "<<y1<<"\n";

include/Layer.h

@@ -2,6 +2,7 @@
 #define LAYER_H
 
 #include<iostream>
+#include<vector>
 #include "utils.h"
 #include "Network.h"
 
@@ -359,6 +360,7 @@ public:
     int classes, num;
     dnnType *mask_h, *mask_d; //anchors
     dnnType *bias_h, *bias_d; //anchors
+    std::vector<std::string> classesNames;
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
     int computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh);

include/Yolo3Detection.h

@@ -52,6 +52,13 @@ class Yolo3Detection {
 
         void update(cv::Mat &frame);
 
+        tk::dnn::Yolo* getYoloLayer(int n=0) {
+            if(n<3)
+                return yolo[n];
+            else 
+                return nullptr;
+        }
+
 };
 
 }}

include/pluginsRT/YoloRT.h

@@ -1,6 +1,8 @@
 #include<cassert>
 #include "../kernels.h"
 
+#define YOLORT_CLASSNAME_W 256
+
 class YoloRT : public IPlugin {
 
 
@@ -16,6 +18,7 @@ public:
         if(yolo != nullptr) {
             memcpy(mask, yolo->mask_h, sizeof(dnnType)*num);
             memcpy(bias, yolo->bias_h, sizeof(dnnType)*num*3*2);
+			classesNames = yolo->classesNames;
         }
 	}
 
@@ -72,7 +75,7 @@ public:
 
 
 	virtual size_t getSerializationSize() override {
-		return 5*sizeof(int) + num*sizeof(dnnType) + num*3*2*sizeof(dnnType);
+		return 5*sizeof(int) + num*sizeof(dnnType) + num*3*2*sizeof(dnnType) + YOLORT_CLASSNAME_W*classes*sizeof(char);
 	}
 
 	virtual void serialize(void* buffer) override {
@@ -86,10 +89,20 @@ public:
     		tk::dnn::writeBUF(buf, mask[i]);
         for(int i=0; i<3*2*num; i++)
     		tk::dnn::writeBUF(buf, bias[i]);
+
+		// 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++) {
+				tk::dnn::writeBUF(buf, tmp[j]);
+			}
+		}
 	}
 
 	int c, h, w;
     int classes, num;
+	std::vector<std::string> classesNames;
 
     dnnType *mask;
     dnnType *bias;

src/NetworkRT.cpp

@@ -461,6 +461,15 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
         for(int i=0; i<3*2*r->num; i++)
             r->bias[i] = readBUF<dnnType>(buf);
 
+		// save classes names
+        r->classesNames.resize(r->classes);
+		for(int i=0; i<r->classes; i++) {
+            char tmp[YOLORT_CLASSNAME_W];
+			for(int j=0; j<YOLORT_CLASSNAME_W; j++)
+				tmp[j] = readBUF<char>(buf);
+            r->classesNames[i] = std::string(tmp);
+		}
+
         yolos[n_yolos++] = r;
         return r;
     } 

src/Yolo.cpp

@@ -25,6 +25,12 @@ Yolo::Yolo(Network *net, int classes, int num, std::string fname_weights) :
         readBinaryFile(fname_weights, 3*num*2, &bias_h, &bias_d, seek);
     }
 
+    // init default classes name
+    classesNames.clear();
+    for(int i=0; i<classes; i++) {
+        classesNames.push_back(std::to_string(i));
+    }
+
     // same
     output_dim.n = input_dim.n;
     output_dim.c = input_dim.c;

src/Yolo3Detection.cpp

@@ -38,6 +38,7 @@ bool Yolo3Detection::init(std::string tensor_path) {
         memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*num);
         memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*3*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);

tests/yolo3_berkeley/yolo3_berkeley.cpp

@@ -14,6 +14,11 @@ int main() {
     tk::dnn::Yolo *yolo [3];
     #include "models/Yolo3.h"
 
+    // fill classes names
+    for(int i=0; i<3; i++) {
+        yolo[i]->classesNames = {"person", "car", "truck", "bus", "motor", "bike", "rider", "traffic light", "traffic sign", "train"};
+    }
+
     // Load input
     dnnType *data;
     dnnType *input_h;

tests/yolo3_flir/yolo3_flir.cpp

@@ -15,6 +15,11 @@ int main() {
     tk::dnn::Yolo *yolo [3];
     #include "models/Yolo3.h"
 
+    // fill classes names
+    for(int i=0; i<3; i++) {
+        yolo[i]->classesNames = {"person", "bike", "car"};
+    }
+
     // Load input
     dnnType *data;
     dnnType *input_h;