From f778e1aa998f894654b24c0ab9ad759c0eb14019 Mon Sep 17 00:00:00 2001 From: Micaela Verucchi Date: Wed, 5 Aug 2020 19:55:10 +0200 Subject: [PATCH] Fixed boxes to float, add conf thresh as param Signed-off-by: Micaela Verucchi --- README.md | 3 ++- demo/config.yaml | 2 +- demo/demo/demo.cpp | 5 ++++- demo/demo/map.cpp | 2 +- include/tkDNN/CenternetDetection.h | 2 +- include/tkDNN/DetectionNN.h | 2 +- include/tkDNN/MobilenetDetection.h | 2 +- include/tkDNN/Yolo3Detection.h | 2 +- src/CenternetDetection.cpp | 3 ++- src/MobilenetDetection.cpp | 3 ++- src/Yolo3Detection.cpp | 11 ++++++----- 11 files changed, 22 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index a1b5b16..d9b5755 100644 --- a/README.md +++ b/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 ``` @@ -204,6 +204,7 @@ where * ``````is the number of classes the network is trained on * `````` 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). * `````` if set to 0 the demo will not show the visualization but save the video into result.mp4 (if n-batches ==1) +* `````` 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 diff --git a/demo/config.yaml b/demo/config.yaml index e6f91a7..31ac599 100644 --- a/demo/config.yaml +++ b/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 diff --git a/demo/demo/demo.cpp b/demo/demo/demo.cpp index 76b451d..9f50d0b 100644 --- a/demo/demo/demo.cpp +++ b/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; diff --git a/demo/demo/map.cpp b/demo/demo/map.cpp index d724db0..356e35a 100644 --- a/demo/demo/map.cpp +++ b/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); diff --git a/include/tkDNN/CenternetDetection.h b/include/tkDNN/CenternetDetection.h index 227cb78..3c8cfbb 100644 --- a/include/tkDNN/CenternetDetection.h +++ b/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); }; diff --git a/include/tkDNN/DetectionNN.h b/include/tkDNN/DetectionNN.h index 030cf8f..ba42834 100644 --- a/include/tkDNN/DetectionNN.h +++ b/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. diff --git a/include/tkDNN/MobilenetDetection.h b/include/tkDNN/MobilenetDetection.h index cabd7eb..9a5fedc 100644 --- a/include/tkDNN/MobilenetDetection.h +++ b/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); }; diff --git a/include/tkDNN/Yolo3Detection.h b/include/tkDNN/Yolo3Detection.h index 6d38514..100a720 100644 --- a/include/tkDNN/Yolo3Detection.h +++ b/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); }; diff --git a/src/CenternetDetection.cpp b/src/CenternetDetection.cpp index 9d8df38..394e24a 100644 --- a/src/CenternetDetection.cpp +++ b/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; diff --git a/src/MobilenetDetection.cpp b/src/MobilenetDetection.cpp index c905fea..3c54e28 100644 --- a/src/MobilenetDetection.cpp +++ b/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]; diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp index 606c6d1..e9b0064 100644 --- a/src/Yolo3Detection.cpp +++ b/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