From c7941666ec65936fae3599452cf150a0669f75bc Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Tue, 19 Feb 2019 08:43:35 +0000
Subject: [PATCH] demo for more yolo3

---
 include/Network.h      |  2 +-
 src/NetworkRT.cpp      |  1 +
 src/Yolo.cpp           | 29 +++++++++++++++++++++++++++++
 src/Yolo3Detection.cpp | 25 ++++++++++++++++---------
 4 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/include/Network.h b/include/Network.h
index dc2b5b3..958129c 100644
--- a/include/Network.h
+++ b/include/Network.h
@@ -15,7 +15,7 @@ namespace tk { namespace dnn {
 */
 struct dataDim_t {
 
-    int n = 0, c = 0, h = 0, w = 0, l = 0;
+    int n, c, h, w, l;
 
     dataDim_t() : n(1), c(1), h(1), w(1), l(1) {};
 
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index 9b78e99..671df8b 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -392,6 +392,7 @@ bool NetworkRT::deserialize(const char *filename) {
         file.close();
     }
 
+    pluginFactory = new PluginFactory();
     runtimeRT = createInferRuntime(loggerRT);
     engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) pluginFactory);
     //if (gieModelStream) delete [] gieModelStream;
diff --git a/src/Yolo.cpp b/src/Yolo.cpp
index fc3d3b9..0dcea7e 100644
--- a/src/Yolo.cpp
+++ b/src/Yolo.cpp
@@ -75,6 +75,34 @@ dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) {
     return dstData;
 }
 
+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;
+    }
+}
+
 int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh) {
 
     if(predictions == nullptr)
@@ -114,6 +142,7 @@ int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int net
         }
     }
 
+    correct_yolo_boxes(dets + ndets, count, netw, neth, netw, neth, 0);
     ndets = count;
     return count;
 }
diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp
index 69547b0..ef61ae2 100644
--- a/src/Yolo3Detection.cpp
+++ b/src/Yolo3Detection.cpp
@@ -2,6 +2,18 @@
 
 namespace tk { namespace dnn {
 
+float _colors[6][3] = { {1,0,1}, {0,0,1},{0,1,1},{0,1,0},{1,1,0},{1,0,0} };
+float get_color(int c, int x, int max)
+{
+    float ratio = ((float)x/max)*5;
+    int i = floor(ratio);
+    int j = ceil(ratio);
+    ratio -= i;
+    float r = (1-ratio) * _colors[i][c] + ratio*_colors[j][c];
+    //printf("%f\n", r);
+    return r;
+}
+
 bool Yolo3Detection::init(std::string tensor_path) {
 
     //const char *tensor_path = "../data/yolo3/yolo3_berkeley.rt";
@@ -34,15 +46,10 @@ bool Yolo3Detection::init(std::string tensor_path) {
 
     // class colors precompute    
     for(int c=0; c<classes; c++) {
-        int cc = c+1;
-        double d = 1.0*( (cc%16)/8 );
-        double r = 1.0*( (cc%8)/4 ) + (0.5*d);
-        double g = 1.0*( (cc%4)/2 ) + (0.5*d);
-        double b = 1.0*( (cc%2)/1 ) + (0.5*d);
-        if(r > 1) r = 1;
-        if(g > 1) g = 1;
-        if(b > 1) b = 1;
-        //std::cout<<r<<" "<<g<<" "<<b<<"\n";
+        int offset = c*123457 % classes;
+        float r = get_color(2, offset, classes);
+        float g = get_color(1, offset, classes);
+        float b = get_color(0, offset, classes);
         colors[c] = cv::Scalar(int(255.0*b), int(255.0*g), int(255.0*r));
     }
     return true;