From ec02c7292fa7113140b923c25535aca17618e670 Mon Sep 17 00:00:00 2001 From: Francesco Gatti Date: Mon, 16 Sep 2019 19:41:59 +0200 Subject: [PATCH] save layer names in rt file --- demo/demo/demo.cpp | 6 +----- include/Layer.h | 2 ++ include/Yolo3Detection.h | 7 +++++++ include/pluginsRT/YoloRT.h | 15 ++++++++++++++- src/NetworkRT.cpp | 9 +++++++++ src/Yolo.cpp | 6 ++++++ src/Yolo3Detection.cpp | 1 + tests/yolo3_berkeley/yolo3_berkeley.cpp | 5 +++++ tests/yolo3_flir/yolo3_flir.cpp | 5 +++++ 9 files changed, 50 insertions(+), 6 deletions(-) diff --git a/demo/demo/demo.cpp b/demo/demo/demo.cpp index 7ed80da..ac6c63d 100644 --- a/demo/demo/demo.cpp +++ b/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< +#include #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 classesNames; virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); int computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh); diff --git a/include/Yolo3Detection.h b/include/Yolo3Detection.h index 0ff4b1e..0a7ba25 100644 --- a/include/Yolo3Detection.h +++ b/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; + } + }; }} diff --git a/include/pluginsRT/YoloRT.h b/include/pluginsRT/YoloRT.h index 2e52142..dab7c18 100644 --- a/include/pluginsRT/YoloRT.h +++ b/include/pluginsRT/YoloRT.h @@ -1,6 +1,8 @@ #include #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 classesNames; dnnType *mask; dnnType *bias; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 2a73efa..7430b6f 100644 --- a/src/NetworkRT.cpp +++ b/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(buf); + // save classes names + r->classesNames.resize(r->classes); + for(int i=0; iclasses; i++) { + char tmp[YOLORT_CLASSNAME_W]; + for(int j=0; j(buf); + r->classesNames[i] = std::string(tmp); + } + yolos[n_yolos++] = r; return r; } diff --git a/src/Yolo.cpp b/src/Yolo.cpp index f2ec53c..babef38 100644 --- a/src/Yolo.cpp +++ b/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; imask_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); diff --git a/tests/yolo3_berkeley/yolo3_berkeley.cpp b/tests/yolo3_berkeley/yolo3_berkeley.cpp index 43f5da5..70e8b60 100644 --- a/tests/yolo3_berkeley/yolo3_berkeley.cpp +++ b/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; diff --git a/tests/yolo3_flir/yolo3_flir.cpp b/tests/yolo3_flir/yolo3_flir.cpp index 2bc45e4..83b53af 100644 --- a/tests/yolo3_flir/yolo3_flir.cpp +++ b/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;