Fixed boxes to float, add conf thresh as param

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>

README.md

@@ -193,7 +193,7 @@ Once you have succesfully created your rt file, run the demo:
 ```
 ./demo yolo4_fp32.rt ../demo/yolo_test.mp4 y
 ```
-In general the demo program takes 6 parameters:
+In general the demo program takes 7 parameters:
 ```
 ./demo <network-rt-file> <path-to-video> <kind-of-network> <number-of-classes> <n-batches> <show-flag>
 ```
@@ -204,6 +204,7 @@ where
 *  ```<number-of-classes>```is the number of classes the network is trained on
 *  ```<n-batches>``` number of batches to use in inference (N.B. you should first export TKDNN_BATCHSIZE to the required n_batches and create again the rt file for the network).
 *  ```<show-flag>``` if set to 0 the demo will not show the visualization but save the video into result.mp4 (if n-batches ==1)
+*  ```<conf-thresh>``` confidence threshold for the detector. Only bounding boxes with threshold greater than conf-thresh will be displayed.
 
 N.b. By default it is used FP32 inference
 

demo/config.yaml

@@ -3,5 +3,5 @@ map_points  : 101   #number of recall points (0 for all, 101 for COCO, 11 Pascal
 map_levels  : 10    #number of IoU step for the AP
 map_step    : 0.05  #step of IoU
 IoU_thresh  : 0.5   #starting IoU threshold
-conf_thresh : 0.0   #threshold on the condifence of the bbox
+conf_thresh : 0.001   #threshold on the condifence of the bbox
 verbose     : false #print on screen information

demo/demo/demo.cpp

@@ -40,6 +40,9 @@ int main(int argc, char *argv[]) {
     bool show = true;
     if(argc > 6)
         show = atoi(argv[6]); 
+    float conf_thresh=0.3;
+    if(argc > 7)
+        conf_thresh = atof(argv[7]);     
 
     if(n_batch < 1 || n_batch > 64)
         FatalError("Batch dim not supported");
@@ -69,7 +72,7 @@ int main(int argc, char *argv[]) {
         FatalError("Network type not allowed (3rd parameter)\n");
     }
 
-    detNN->init(net, n_classes, n_batch);
+    detNN->init(net, n_classes, n_batch, conf_thresh);
 
     gRun = true;
 

demo/demo/map.cpp

@@ -105,7 +105,7 @@ int main(int argc, char *argv[])
         default:
             FatalError("Network type not allowed (3rd parameter)\n");
     }
-    detNN->init(net, n_classes);
+    detNN->init(net, n_classes, 1, conf_thresh);
 
     //read images 
     std::ifstream all_labels(labels_path);

include/tkDNN/CenternetDetection.h

@@ -73,7 +73,7 @@ public:
     CenternetDetection() {};
     ~CenternetDetection() {}; 
 
-    bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1);
+    bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
     void preprocess(cv::Mat &frame, const int bi=0);
     void postprocess(const int bi=0,const bool mAP=false);
 };

include/tkDNN/DetectionNN.h

@@ -84,7 +84,7 @@ class DetectionNN {
          * @param n_batches maximum number of batches to use in inference
          * @return true if everything is correct, false otherwise.
          */
-        virtual bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1) = 0;
+        virtual bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3) = 0;
         
         /**
          * This method performs the whole detection of the NN.

include/tkDNN/MobilenetDetection.h

@@ -65,7 +65,7 @@ public:
     MobilenetDetection() {};
     ~MobilenetDetection() {}; 
 
-    bool init(const std::string& tensor_path, const int n_classes, const int n_batches=1);
+    bool init(const std::string& tensor_path, const int n_classes, const int n_batches=1, const float conf_thresh=0.3);
     void preprocess(cv::Mat &frame, const int bi=0);
     void postprocess(const int bi=0,const bool mAP=false);
 };

include/tkDNN/Yolo3Detection.h

@@ -24,7 +24,7 @@ public:
     Yolo3Detection() {};
     ~Yolo3Detection() {}; 
 
-    bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1);
+    bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
     void preprocess(cv::Mat &frame, const int bi=0);
     void postprocess(const int bi=0,const bool mAP=false);
 };

src/CenternetDetection.cpp

@@ -3,11 +3,12 @@
 
 namespace tk { namespace dnn {
 
-bool CenternetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches){
+bool CenternetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh){
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
     classes = n_classes;
     nBatches = n_batches;
+    confThreshold = conf_thresh;
 
     dim = netRT->input_dim;
 

src/MobilenetDetection.cpp

@@ -126,12 +126,13 @@ float MobilenetDetection::iou(const tk::dnn::box &a, const tk::dnn::box &b){
     return iou;
 }
 
-bool MobilenetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches){
+bool MobilenetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh){
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
     imageSize = netRT->input_dim.h;
     classes = n_classes;
     nBatches = n_batches;
+    confThreshold = conf_thresh;
 
     SSDSpec specs[N_SSDSPEC];
 

src/Yolo3Detection.cpp

@@ -3,13 +3,14 @@
 
 namespace tk { namespace dnn {
 
-bool Yolo3Detection::init(const std::string& tensor_path, const int n_classes, const int n_batches) {
+bool Yolo3Detection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh) {
 
     //convert network to tensorRT
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
 
     nBatches = n_batches;
+    confThreshold = conf_thresh;
     tk::dnn::dataDim_t idim = netRT->input_dim;    
     idim.n = nBatches;
 
@@ -109,10 +110,10 @@ void Yolo3Detection::postprocess(const int bi, const bool mAP){
     detected.clear();
     for(int j=0; j<nDets; j++) {
         tk::dnn::Yolo::box b = dets[j].bbox;
-        int x0   = (b.x-b.w/2.);
+        float x0   = (b.x-b.w/2.);
-        int x1   = (b.x+b.w/2.);
+        float x1   = (b.x+b.w/2.);
-        int y0   = (b.y-b.h/2.);
+        float y0   = (b.y-b.h/2.);
-        int y1   = (b.y+b.h/2.);
+        float y1   = (b.y+b.h/2.);
 
         // convert to image coords
         x0 = x_ratio*x0;