diff --git a/demo/demo/demo.cpp b/demo/demo/demo.cpp index 76b451d..85538a8 100644 --- a/demo/demo/demo.cpp +++ b/demo/demo/demo.cpp @@ -69,7 +69,8 @@ int main(int argc, char *argv[]) { FatalError("Network type not allowed (3rd parameter)\n"); } - detNN->init(net, n_classes, n_batch); + float conf_thresh = 0.3; + detNN->init(net, conf_thresh, n_classes, n_batch); gRun = true; diff --git a/demo/demo/map.cpp b/demo/demo/map.cpp index d724db0..ea6eba4 100644 --- a/demo/demo/map.cpp +++ b/demo/demo/map.cpp @@ -31,6 +31,8 @@ int main(int argc, char *argv[]) const char *config_filename = "../demo/config.yaml"; const char * net = "yolo3.rt"; const char * labels_path = "../demo/COCO_val2017/all_labels.txt"; + int n_batches = 1; + float confidence_thresh = 0.3; bool show = false; bool write_dets = false; bool write_res_on_file = true; @@ -53,6 +55,12 @@ int main(int argc, char *argv[]) labels_path = argv[3]; if(argc > 4) config_filename = argv[4]; + if(argc > 5) + n_batches = atoi(argv[5]); + if(argc > 6) + confidence_thresh = atof(argv[6]); + + std::cout<<"conf t: "<init(net, n_classes); + detNN->init(net, confidence_thresh, n_classes, n_batches); //read images std::ifstream all_labels(labels_path); @@ -118,90 +126,109 @@ int main(int argc, char *argv[]) if(show) cv::namedWindow("detection", cv::WINDOW_NORMAL); + bool file_ok = false; + int images_done; - for (images_done=0 ; std::getline(all_labels, l_filename) && images_done < n_images ; ++images_done) { - std::cout < batch_frames; - batch_frames.push_back(frame); - int height = frame.rows; - int width = frame.cols; - - if(!frame.data) - break; std::vector batch_dnn_input; - batch_dnn_input.push_back(frame.clone()); + + std::vector cur_frames; + for(;cur_batches> b.cl >> b.x >> b.y >> b.w >> b.h; + b.prob = 1; + b.truthFlag = 1; + f.gt.push_back(b); + + if(show)// draw rectangle for groundtruth + cv::rectangle(batch_frames[cur_batches], cv::Point((b.x-b.w/2)*f.width, (b.y-b.h/2)*f.height), cv::Point((b.x+b.w/2)*f.width,(b.y+b.h/2)*f.height), cv::Scalar(0, 255, 0), 2); + } + } + + cur_frames.push_back(f); + } + if (!file_ok) + break; //inference - detected_bbox.clear(); - detNN->update(batch_dnn_input,1,write_res_on_file, ×, write_coco_json); + detNN->update(batch_dnn_input,cur_batches,write_res_on_file, ×, write_coco_json); detNN->draw(batch_frames); - detected_bbox = detNN->detected; - if(write_coco_json) - printJsonCOCOFormat(&coco_json, f.iFilename.c_str(), detected_bbox, classes, width, height); + for(int j=0;jbatchDetected[j], classes, cur_frames[j].width, cur_frames[j].height); - std::ofstream myfile; - if(write_dets) - myfile.open ("det/"+f.lFilename.substr(f.lFilename.find("labels/") + 7)); + std::ofstream myfile; + if(write_dets) + myfile.open ("det/"+cur_frames[j].lFilename.substr(cur_frames[j].lFilename.find("labels/") + 7)); - // save detections labels - for(auto d:detected_bbox){ - //convert detected bb in the same format as label - /// / / / - tk::dnn::BoundingBox b; - b.x = (d.x + d.w/2) / width; - b.y = (d.y + d.h/2) / height; - b.w = d.w / width; - b.h = d.h / height; - b.prob = d.prob; - b.cl = d.cl; - f.det.push_back(b); + // save detections labels + for(auto d:detNN->batchDetected[j]){ + //convert detected bb in the same format as label + /// / / / + tk::dnn::BoundingBox b; + b.x = (d.x + d.w/2) / cur_frames[j].width; + b.y = (d.y + d.h/2) / cur_frames[j].height; + b.w = d.w / cur_frames[j].width; + b.h = d.h / cur_frames[j].height; + b.prob = d.prob; + b.cl = d.cl; + cur_frames[j].det.push_back(b); + + if(write_dets) + myfile << d.cl << " "<< d.prob << " "<< b.x << " "<< b.y << " "<< b.w << " "<< b.h <<"\n"; + + if(show)// draw rectangle for detection + cv::rectangle(batch_frames[j], cv::Point(d.x, d.y), cv::Point(d.x + d.w, d.y + d.h), cv::Scalar(0, 0, 255), 2); + } if(write_dets) - myfile << d.cl << " "<< d.prob << " "<< b.x << " "<< b.y << " "<< b.w << " "<< b.h <<"\n"; - - if(show)// draw rectangle for detection - cv::rectangle(batch_frames[0], cv::Point(d.x, d.y), cv::Point(d.x + d.w, d.y + d.h), cv::Scalar(0, 0, 255), 2); - } - - if(write_dets) - myfile.close(); - - // read and save groundtruth labels - if(fileExist(f.lFilename.c_str())) - { - std::ifstream labels(l_filename); - for(std::string line; std::getline(labels, line); ){ - std::istringstream in(line); - tk::dnn::BoundingBox b; - in >> b.cl >> b.x >> b.y >> b.w >> b.h; - b.prob = 1; - b.truthFlag = 1; - f.gt.push_back(b); - - if(show)// draw rectangle for groundtruth - cv::rectangle(batch_frames[0], cv::Point((b.x-b.w/2)*width, (b.y-b.h/2)*height), cv::Point((b.x+b.w/2)*width,(b.y+b.h/2)*height), cv::Scalar(0, 255, 0), 2); - } - } + myfile.close(); - images.push_back(f); + images.push_back(cur_frames[j]); - if(show){ - cv::imshow("detection", batch_frames[0]); - cv::waitKey(0); + if(show){ + cv::imshow("detection", batch_frames[j]); + cv::waitKey(0); + } + } - + std::cout < gt; std::vector det; + int width; + int height; void print() const; }; diff --git a/src/CenternetDetection.cpp b/src/CenternetDetection.cpp index 9d8df38..9c3436c 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 float conf_threshold, const int n_classes, const int n_batches){ 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_threshold; dim = netRT->input_dim; diff --git a/src/MobilenetDetection.cpp b/src/MobilenetDetection.cpp index c905fea..af20f9e 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 float conf_threshold, const int n_classes, const int n_batches){ 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_threshold; SSDSpec specs[N_SSDSPEC]; diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp index c76af20..09eb151 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 float conf_threshold, const int n_classes, const int n_batches) { //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_threshold; tk::dnn::dataDim_t idim = netRT->input_dim; idim.n = nBatches;