yoloRT load anchors

demo/demo/demo.cpp

@@ -24,7 +24,7 @@ int main(int argc, char *argv[]) {
     signal(SIGINT, sig_handler);
 
     tk::dnn::Yolo3Detection yolo;
-    yolo.init("yolo3_berkeley");
+    yolo.init("yolo3_berkeley.rt");
 
     gRun = true;
 

include/Layer.h

@@ -343,7 +343,7 @@ public:
         float x, y, w, h;
     };
 
-    typedef struct detection{
+    struct detection{
         Yolo::box bbox;
         int classes;
         float *prob;
@@ -361,7 +361,7 @@ public:
     dnnType *bias_h, *bias_d; //anchors
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
-    int computeDetections(Yolo::detection *dets, int &ndets, int w, int h, int netw, int neth, float thresh);
+    int computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh);
 
     dnnType *predictions;
 

include/Network.h

@@ -15,7 +15,7 @@ namespace tk { namespace dnn {
 */
 struct dataDim_t {
 
-    int n, c, h, w, l;
+    int n = 0, c = 0, h = 0, w = 0, l = 0;
 
     dataDim_t() : n(1), c(1), h(1), w(1), l(1) {};
 

include/NetworkRT.h

@@ -1,6 +1,7 @@
 #ifndef NETWORKRT_H
 #define NETWORKRT_H
 
+#include <string.h> // memcpy
 #include "utils.h"
 #include "Network.h"
 #include "Layer.h"
@@ -8,6 +9,40 @@
 
 namespace tk { namespace dnn {
 
+template<typename T> void writeBUF(char*& buffer, const T& val)
+{
+    *reinterpret_cast<T*>(buffer) = val;
+    buffer += sizeof(T);
+}
+
+template<typename T> T readBUF(const char*& buffer)
+{
+    T val = *reinterpret_cast<const T*>(buffer);
+    buffer += sizeof(T);
+    return val;
+}
+
+using namespace nvinfer1;
+#include "pluginsRT/ActivationLeakyRT.h"
+#include "pluginsRT/ReorgRT.h"
+#include "pluginsRT/RegionRT.h"
+//#include "pluginsRT/RouteRT.h"
+#include "pluginsRT/ShortcutRT.h"
+#include "pluginsRT/YoloRT.h"
+#include "pluginsRT/UpsampleRT.h"
+//#include "pluginsRT/Int8Calibrator.h"
+
+class PluginFactory : IPluginFactory
+{
+public:
+    YoloRT *yolos[16];
+    int n_yolos;
+
+	virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength);
+};
+
+
+
 class NetworkRT {
 
 public:
@@ -27,6 +62,8 @@ public:
     dnnType *output;
     cudaStream_t stream;
 
+    PluginFactory *pluginFactory;
+
     NetworkRT(Network *net, const char *name);
     virtual ~NetworkRT();
 
@@ -53,19 +90,5 @@ public:
     bool deserialize(const char *filename);
 };
 
-
-template<typename T> void writeBUF(char*& buffer, const T& val)
-{
-    *reinterpret_cast<T*>(buffer) = val;
-    buffer += sizeof(T);
-}
-
-template<typename T> T readBUF(const char*& buffer)
-{
-    T val = *reinterpret_cast<const T*>(buffer);
-    buffer += sizeof(T);
-    return val;
-}
-
 }}
 #endif //NETWORKRT_H

include/Yolo3Detection.h

@@ -31,8 +31,8 @@ class Yolo3Detection {
         cv::Mat bgr[3]; 
 
     public:
-        int classes = 10;
+        int classes = 0;
-        int num = 3;
+        int num = 0;
         float thresh = 0.3;
         cv::Scalar colors[256];
 

src/NetworkRT.cpp

@@ -11,14 +11,6 @@
 #include "NetworkRT.h"
 
 using namespace nvinfer1;
-#include "pluginsRT/ActivationLeakyRT.cpp"
-#include "pluginsRT/ReorgRT.cpp"
-#include "pluginsRT/RegionRT.cpp"
-//#include "pluginsRT/RouteRT.cpp"
-#include "pluginsRT/ShortcutRT.cpp"
-#include "pluginsRT/YoloRT.cpp"
-#include "pluginsRT/UpsampleRT.cpp"
-#include "pluginsRT/Int8Calibrator.cpp"
 
 // Logger for info/warning/errors
 class Logger : public ILogger {
@@ -54,14 +46,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
         builderRT->setMaxBatchSize(1);
         builderRT->setMaxWorkspaceSize(1 << 30);
 
-        //change datatype based on system specs
+        if(net->fp16 && builderRT->platformHasFastFp16()) {
-        if(builderRT->platformHasFastInt8()) {
-            BatchStream bstream({32,dim.c, dim.h, dim.w}, 32, 1);
-            Int8EntropyCalibrator calib(bstream, 0, false);
-            builderRT->setInt8Mode(true);
-            builderRT->setInt8Calibrator(&calib);
-
-        } else if(net->fp16 && builderRT->platformHasFastFp16()) {
             dtRT = DataType::kHALF;
             builderRT->setHalf2Mode(true);
         }
@@ -393,87 +378,6 @@ bool NetworkRT::serialize(const char *filename) {
     return true;
 }
 
-class PluginFactory : IPluginFactory
-{
-public:
-
-	virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) {
-        const char * buf = reinterpret_cast<const char*>(serialData);
-        
-        std::string name(layerName);
-
-        if(name.find("Activation") == 0) {
-            ActivationLeakyRT *a = new ActivationLeakyRT();
-            a->size = readBUF<int>(buf);
-            return a;
-        }
-
-        if(name.find("Region") == 0) {
-            RegionRT *r = new RegionRT(readBUF<int>(buf),    //classes
-                                       readBUF<int>(buf),    //coords
-                                       readBUF<int>(buf));   //num
-
-        	r->c = readBUF<int>(buf);
-		    r->h = readBUF<int>(buf);
-		    r->w = readBUF<int>(buf);
-            return r;
-        } 
-
-        if(name.find("Reorg") == 0) {
-            ReorgRT *r = new ReorgRT(readBUF<int>(buf)); //stride
-        	r->c = readBUF<int>(buf);
-		    r->h = readBUF<int>(buf);
-		    r->w = readBUF<int>(buf);
-            return r;
-        } 
-
-        if(name.find("Shortcut") == 0) {
-            ShortcutRT *r = new ShortcutRT();
-        	r->c = readBUF<int>(buf);
-		    r->h = readBUF<int>(buf);
-		    r->w = readBUF<int>(buf);
-            return r;
-        } 
-
-        if(name.find("Yolo") == 0) {
-            YoloRT *r = new YoloRT(readBUF<int>(buf),    //classes
-                                   readBUF<int>(buf));   //num
-        	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;
-        } 
-
-        if(name.find("Upsample") == 0) {
-            UpsampleRT *r = new UpsampleRT(readBUF<int>(buf)); //stride
-        	r->c = readBUF<int>(buf);
-		    r->h = readBUF<int>(buf);
-		    r->w = readBUF<int>(buf);
-            return r;
-        } 
-/*
-        if(name.find("Route") == 0) {
-            RouteRT *r = new RouteRT();
-            r->in = readBUF<int>(buf);
-            for(int i=0; i<RouteRT::MAX_INPUTS; i++)
-                r->c_in[i] = readBUF<int>(buf);
-        	r->c = readBUF<int>(buf);
-		    r->h = readBUF<int>(buf);
-		    r->w = readBUF<int>(buf);
-            return r;
-        } 
-*/
-        FatalError("Cant deserialize Plugin");
-        return NULL;
-    }
-};
-
 bool NetworkRT::deserialize(const char *filename) {
 
     char *gieModelStream{nullptr};
@@ -488,13 +392,89 @@ bool NetworkRT::deserialize(const char *filename) {
         file.close();
     }
 
-    PluginFactory plfact;
-
     runtimeRT = createInferRuntime(loggerRT);
-    engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) &plfact);
+    engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) pluginFactory);
     //if (gieModelStream) delete [] gieModelStream;
 
     return true;
 }
 
+
+
+IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) {
+    const char * buf = reinterpret_cast<const char*>(serialData);
+    
+    std::string name(layerName);
+
+    if(name.find("Activation") == 0) {
+        ActivationLeakyRT *a = new ActivationLeakyRT();
+        a->size = readBUF<int>(buf);
+        return a;
+    }
+
+    if(name.find("Region") == 0) {
+        RegionRT *r = new RegionRT(readBUF<int>(buf),    //classes
+                                    readBUF<int>(buf),    //coords
+                                    readBUF<int>(buf));   //num
+
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        return r;
+    } 
+
+    if(name.find("Reorg") == 0) {
+        ReorgRT *r = new ReorgRT(readBUF<int>(buf)); //stride
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        return r;
+    } 
+
+    if(name.find("Shortcut") == 0) {
+        ShortcutRT *r = new ShortcutRT();
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        return r;
+    } 
+
+    if(name.find("Yolo") == 0) {
+        YoloRT *r = new YoloRT(readBUF<int>(buf),    //classes
+                                readBUF<int>(buf));   //num
+        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);
+
+        yolos[n_yolos++] = r;
+        return r;
+    } 
+
+    if(name.find("Upsample") == 0) {
+        UpsampleRT *r = new UpsampleRT(readBUF<int>(buf)); //stride
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        return r;
+    } 
+/*
+    if(name.find("Route") == 0) {
+        RouteRT *r = new RouteRT();
+        r->in = readBUF<int>(buf);
+        for(int i=0; i<RouteRT::MAX_INPUTS; i++)
+            r->c_in[i] = readBUF<int>(buf);
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        return r;
+    } 
+*/
+    FatalError("Cant deserialize Plugin");
+    return NULL;
+}
+
 }}

src/Yolo.cpp

@@ -18,13 +18,12 @@ Yolo::Yolo(Network *net, int classes, int num, const char* fname_weights) :
     this->num = num;
 
     // load anchors
-    int seek = 0;
+    if(fname_weights != nullptr) {
-    readBinaryFile(fname_weights, num, &mask_h, &mask_d, seek);
+        int seek = 0;
-    seek += num;
+        readBinaryFile(fname_weights, num, &mask_h, &mask_d, seek);
-    readBinaryFile(fname_weights, 3*num*2, &bias_h, &bias_d, seek);
+        seek += num;
-
+        readBinaryFile(fname_weights, 3*num*2, &bias_h, &bias_d, seek);
-    printDeviceVector(num, mask_h, false);
+    }
-    printDeviceVector(3*num*2, bias_h, false);
 
     // same
     output_dim.n = input_dim.n;
@@ -33,10 +32,6 @@ Yolo::Yolo(Network *net, int classes, int num, const char* fname_weights) :
     output_dim.w = input_dim.w;
     output_dim.l = input_dim.l;
 
-    std::cout<<"YOLO INPUT: ";
-    input_dim.print();
-    std::cout<<"\n";
-
     checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
     predictions = nullptr;
 }
@@ -62,35 +57,6 @@ Yolo::box get_yolo_box(float *x, float *biases, int n, int index, int i, int j,
     return b;
 }
 
-void correct_yolo_boxes(Yolo::detection *dets, int n, int w, int h, int netw, int neth, int relative)
-{
-    int i;
-    int new_w=0;
-    int new_h=0;
-    if (((float)netw/w) < ((float)neth/h)) {
-        new_w = netw;
-        new_h = (h * netw)/w;
-    } else {
-        new_h = neth;
-        new_w = (w * neth)/h;
-    }
-    for (i = 0; i < n; ++i){
-        Yolo::box b = dets[i].bbox;
-        b.x =  (b.x - (netw - new_w)/2./netw) / ((float)new_w/netw); 
-        b.y =  (b.y - (neth - new_h)/2./neth) / ((float)new_h/neth); 
-        b.w *= (float)netw/new_w;
-        b.h *= (float)neth/new_h;
-        if(!relative){
-            b.x *= w;
-            b.w *= w;
-            b.y *= h;
-            b.h *= h;
-        }
-        dets[i].bbox = b;
-    }
-}
-
-
 dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) {
 
     checkCuda( cudaMemcpy(dstData, srcData, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice));
@@ -109,14 +75,12 @@ dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) {
     return dstData;
 }
 
-int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int w, int h, int netw, int neth, float thresh) {
+int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh) {
 
     if(predictions == nullptr)
         predictions = new dnnType[output_dim.tot()];
     checkCuda( cudaMemcpy(predictions, dstData, output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToHost));
 
-    int relative = 0;
-
     int lw = output_dim.w;
     int lh = output_dim.h;
 
@@ -150,7 +114,6 @@ int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int w, int h, int
         }
     }
 
-    correct_yolo_boxes(dets + ndets, count, w, h, netw, neth, relative);
     ndets = count;
     return count;
 }

src/Yolo3Detection.cpp

@@ -2,10 +2,36 @@
 
 namespace tk { namespace dnn {
 
-bool Yolo3Detection::init(std::string tensor_folder) {
+bool Yolo3Detection::init(std::string tensor_path) {
 
     //const char *tensor_path = "../data/yolo3/yolo3_berkeley.rt";
 
+    //convert network to tensorRT
+    std::cout<<(tensor_path).c_str()<<"\n";
+    netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
+
+    if(netRT->pluginFactory->n_yolos != 3) {
+        FatalError("this is not yolo3");
+    }
+
+    for(int i=0; i<netRT->pluginFactory->n_yolos; i++) {
+        YoloRT *yRT = netRT->pluginFactory->yolos[i];
+        classes = yRT->classes;
+        num = yRT->num;
+
+        // make a yolo layer for interpret predictions
+        yolo[i] = new tk::dnn::Yolo(nullptr, classes, num, nullptr); // yolo without input and bias
+        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);
+    }
+
+    dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
+    
+    checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot()));
+    checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot()));
+
+
     // class colors precompute    
     for(int c=0; c<classes; c++) {
         int cc = c+1;
@@ -19,23 +45,6 @@ bool Yolo3Detection::init(std::string tensor_folder) {
         //std::cout<<r<<" "<<g<<" "<<b<<"\n";
         colors[c] = cv::Scalar(int(255.0*b), int(255.0*g), int(255.0*r));
     }
-
-    //convert network to tensorRT
-    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 + "_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 + "_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 + "_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);
-    
-    checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot()));
-    checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot()));
-
     return true;
 } 
 
@@ -82,7 +91,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
     for(int i=0; i<3; i++) {
         rt_out[i] = (dnnType*)netRT->buffersRT[i+1];
         yolo[i]->dstData = rt_out[i];
-        yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, netRT->input_dim.w, netRT->input_dim.h, thresh);
+        yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, thresh);
     }
     tk::dnn::Yolo::mergeDetections(dets, ndets, classes);
     TIMER_STOP

tests/yolo3_berkeley/yolo3_berkeley.cpp

@@ -326,9 +326,9 @@ int main() {
     int ndets = 0;
     int classes = yolo0.classes;
     tk::dnn::Yolo::detection *dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
-    yolo0.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5);
+    yolo0.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5);
-    yolo1.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5);
+    yolo1.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5);
-    yolo2.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5);
+    yolo2.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5);
     tk::dnn::Yolo::mergeDetections(dets, ndets, classes);
 
     for(int j=0; j<ndets; j++) {
@@ -369,14 +369,5 @@ int main() {
         std::cout<<"TRT   vs correct"; checkResult(odim, rt_out[i], out);
         std::cout<<"CUDNN vs TRT    "; checkResult(odim, cudnn_out[i], rt_out[i]);
     }
-
-    std::cout<<"copyng layer config to this folder\n";
-    std::string cmd;
-    cmd = "cp " + std::string(g82_bin) + " yolo3_berkeley_0.bin";  
-    std::cout<<cmd<<"\n"; system(cmd.c_str());
-    cmd = "cp " + std::string(g94_bin) + " yolo3_berkeley_1.bin";  
-    std::cout<<cmd<<"\n"; system(cmd.c_str());
-    cmd = "cp " + std::string(g106_bin) + " yolo3_berkeley_2.bin";  
-    std::cout<<cmd<<"\n"; system(cmd.c_str());
     return 0;
 }

tests/yolo3_voc/yolo3_voc.cpp

@@ -326,9 +326,9 @@ int main() {
     int ndets = 0;
     int classes = yolo0.classes;
     tk::dnn::Yolo::detection *dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
-    yolo0.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5);
+    yolo0.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5);
-    yolo1.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5);
+    yolo1.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5);
-    yolo2.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5);
+    yolo2.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5);
     tk::dnn::Yolo::mergeDetections(dets, ndets, classes);
 
     for(int j=0; j<ndets; j++) {
@@ -369,14 +369,5 @@ int main() {
         std::cout<<"TRT   vs correct"; checkResult(odim, rt_out[i], out);
         std::cout<<"CUDNN vs TRT    "; checkResult(odim, cudnn_out[i], rt_out[i]);
     }
-
-    std::cout<<"copyng layer config to this folder\n";
-    std::string cmd;
-    cmd = "cp " + std::string(g82_bin) + " yolo3_voc_0.bin";  
-    std::cout<<cmd<<"\n"; system(cmd.c_str());
-    cmd = "cp " + std::string(g94_bin) + " yolo3_voc_1.bin";  
-    std::cout<<cmd<<"\n"; system(cmd.c_str());
-    cmd = "cp " + std::string(g106_bin) + " yolo3_voc_2.bin";  
-    std::cout<<cmd<<"\n"; system(cmd.c_str());
     return 0;
 }