mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Add variable batch for detector update

Browse Source 
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
2020-05-19 09:32:27 +02:00
parent ad21cc0b79
commit f2991bec5c
3 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo/demo/demo.cpp
+1 -1
View File
@@ -111,7 +111,7 @@ int main(int argc, char *argv[]) {
break; break;
//inference //inference
detNN->update(batch_dnn_input); detNN->update(batch_dnn_input, n_batch);
detNN->draw(batch_frame); detNN->draw(batch_frame);
if(show){ if(show){
demo/demo/map.cpp
+1 -1
View File
@@ -144,7 +144,7 @@ int main(int argc, char *argv[])
//inference //inference
detected_bbox.clear(); detected_bbox.clear();
detNN->update(batch_dnn_input, write_res_on_file, &times, write_coco_json); detNN->update(batch_dnn_input,1,write_res_on_file, &times, write_coco_json);
detNN->draw(batch_frames); detNN->draw(batch_frames);
detected_bbox = detNN->detected; detected_bbox = detNN->detected;
include/tkDNN/DetectionNN.h
+8 -5
View File
@@ -81,7 +81,7 @@ class DetectionNN {
* *
* @param tensor_path path to the rt file og the NN. * @param tensor_path path to the rt file og the NN.
* @param n_classes number of classes for the given dataset. * @param n_classes number of classes for the given dataset.
* @param n_batches number of batches to use in inference * @param n_batches maximum number of batches to use in inference
* @return true if everything is correct, false otherwise. * @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) = 0;
@@ -90,20 +90,23 @@ class DetectionNN {
* This method performs the whole detection of the NN. * This method performs the whole detection of the NN.
* *
* @param frames frames to run detection on. * @param frames frames to run detection on.
* @param cur_batches number of batches to use in inference
* @param save_times if set to true, preprocess, inference and postprocess times * @param save_times if set to true, preprocess, inference and postprocess times
* are saved on a csv file, otherwise not. * are saved on a csv file, otherwise not.
* @param times pointer to the output stream where to write times * @param times pointer to the output stream where to write times
* @param mAP set to true only if all the probabilities for a bounding * @param mAP set to true only if all the probabilities for a bounding
* box are needed, as in some cases for the mAP calculation * box are needed, as in some cases for the mAP calculation
*/ */
void update(std::vector<cv::Mat>& frames, bool save_times=false, std::ofstream *times=nullptr, const bool mAP=false){ void update(std::vector<cv::Mat>& frames, const int cur_batches=1, bool save_times=false, std::ofstream *times=nullptr, const bool mAP=false){
if(save_times && times==nullptr) if(save_times && times==nullptr)
FatalError("save_times set to true, but no valid ofstream given"); FatalError("save_times set to true, but no valid ofstream given");
if(cur_batches > nBatches)
FatalError("A batch size greater than nBatches cannot be used");
if(VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30); if(VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
{ {
TIMER_START TIMER_START
for(int bi=0; bi<nBatches;++bi){ for(int bi=0; bi<cur_batches;++bi){
if(!frames[bi].data) if(!frames[bi].data)
FatalError("No image data feed to detection"); FatalError("No image data feed to detection");
originalSize = frames[bi].size(); originalSize = frames[bi].size();
@@ -115,7 +118,7 @@ class DetectionNN {
//do inference //do inference
tk::dnn::dataDim_t dim = netRT->input_dim; tk::dnn::dataDim_t dim = netRT->input_dim;
dim.n = nBatches; dim.n = cur_batches;
{ {
if(VERBOSE) dim.print(); if(VERBOSE) dim.print();
TIMER_START TIMER_START
@@ -129,7 +132,7 @@ class DetectionNN {
batchDetected.clear(); batchDetected.clear();
{ {
TIMER_START TIMER_START
for(int bi=0; bi<nBatches;++bi) for(int bi=0; bi<cur_batches;++bi)
postprocess(bi, mAP); postprocess(bi, mAP);
TIMER_STOP TIMER_STOP
if(save_times) *times<<t_ns<<"\n"; if(save_times) *times<<t_ns<<"\n";