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

Batch size > 1 for the 3D demo.

Browse Source 
This commit lets to use differtent batch size for 3D CenterNet
and CenterTrack.

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
This commit is contained in:
Davide Sapienza
2021-04-29 11:13:24 +02:00
parent 28fab9c3e1
commit be6ad27c11
6 changed files with 309 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/CenternetDetection3D.h
+4 -8
View File
@@ -74,22 +74,18 @@ private:
int *ids_out;
struct threshold op;
float peakThreshold = 0.2;
float centerThreshold = 0.3; //default 0.5
cv::Mat corners, pts3DHomo;
std::vector<box3D> detected3D;
std::vector<int>cls3D;
std::vector<std::vector<int>> face_id;
public:
CenternetDetection3D() {};
~CenternetDetection3D() {};
bool init(const std::string& tensor_path, const int n_classes=3);
void preprocess(cv::Mat &frame);
void postprocess();
cv::Mat draw(cv::Mat &frame);
bool init(const std::string& tensor_path, const int n_classes=3, 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);
void draw(std::vector<cv::Mat>& frames);
};
include/tkDNN/CenternetDetection3DTrack.h
+7 -6
View File
@@ -145,6 +145,7 @@ private:
int count_det;
//tracks
std::vector<struct trackingRes> tr_res;
std::vector<std::vector<struct trackingRes>> batchTracked;
int count_tr;
int track_id=0;
@@ -153,7 +154,7 @@ private:
bool init_pre_inf();
bool init_postprocessing();
bool init_visualization(const int n_classes);
void pre_inf();
void pre_inf(const int bi);
void _get_additional_inputs();
cv::Mat transform_preds_with_trans(float x1, float x2);
void tracking();
@@ -161,11 +162,11 @@ private:
public:
tk::dnn::Network *pre_phase_net = nullptr;
CenternetDetection3DTrack() {};
~CenternetDetection3DTrack() {};
bool init(const std::string& tensor_path, const int n_classes=3);
void preprocess(cv::Mat &frame);
void postprocess();
cv::Mat draw(cv::Mat &frame);
~CenternetDetection3DTrack() {};
bool init(const std::string& tensor_path, const int n_classes=3, 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);
void draw(std::vector<cv::Mat>& frames);
};
include/tkDNN/DetectionNN3D.h
+52 -31
View File
@@ -3,8 +3,11 @@
#include <iostream>
#include <signal.h>
#include <stdlib.h>
#include <stdlib.h>
#ifdef __linux__
#include <unistd.h>
#endif
#include <mutex>
#include "utils.h"
@@ -30,10 +33,12 @@ class DetectionNN3D {
tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d;
cv::Size originalSize;
std::vector<cv::Size> originalSize;
cv::Scalar colors[256];
int nBatches = 1;
#ifdef OPENCV_CUDACONTRIB
cv::cuda::GpuMat bgr[3];
cv::cuda::GpuMat imagePreproc;
@@ -47,21 +52,26 @@ class DetectionNN3D {
* This method preprocess the image, before feeding it to the NN.
*
* @param frame original frame to adapt for inference.
* @param bi batch index
*/
virtual void preprocess(cv::Mat &frame) = 0;
virtual void preprocess(cv::Mat &frame, const int bi=0) = 0;
/**
* This method postprocess the output of the NN to obtain the correct
* boundig boxes.
*
* @param bi batch index
* @param mAP set to true only if all the probabilities for a bounding
* box are needed, as in some cases for the mAP calculation
*/
virtual void postprocess() = 0;
virtual void postprocess(const int bi=0,const bool mAP=false) = 0;
public:
int classes = 0;
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
std::vector<tk::dnn::box> detected; /*bounding boxes in output*/
std::vector<tk::dnn::box3D> detected3D; /*bounding boxes in output*/
std::vector<std::vector<tk::dnn::box3D>> batchDetected; /*bounding boxes in output*/
std::vector<double> pre_stats, stats, post_stats, visual_stats; /*keeps track of inference times (ms)*/
std::vector<std::string> classesNames;
@@ -69,68 +79,79 @@ class DetectionNN3D {
~DetectionNN3D(){};
/**
* Method used to inialize the class, allocate memory and compute
* Method used to initialize the class, allocate memory and compute
* needed data.
*
* @param tensor_path path to the rt file og the NN.
* @param tensor_path path to the rt file of the NN.
* @param n_classes number of classes for the given dataset.
* @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=3) = 0;
/**
* Method to draw boundixg boxes and labels on a frame.
*
* @param frame orginal frame to draw bounding box on.
* @return frame with boundig boxes.
*/
virtual cv::Mat draw(cv::Mat &frame){};
virtual bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1, const float conf_thresh=0.3) = 0;
/**
* This method performs the whole detection of the NN.
*
* @param frame frame 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
* 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
* box are needed, as in some cases for the mAP calculation.
*/
void update(cv::Mat &frame, bool save_times=false, std::ofstream *times=nullptr){
if(!frame.data)
FatalError("No image data feed to detection");
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)
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");
originalSize = frame.size();
printCenteredTitle(" TENSORRT detection ", '=', 30);
originalSize.clear();
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
{
TKDNN_TSTART
preprocess(frame);
for(int bi=0; bi<cur_batches;++bi){
if(!frames[bi].data)
FatalError("No image data feed to detection");
originalSize.push_back(frames[bi].size());
preprocess(frames[bi], bi);
}
TKDNN_TSTOP
pre_stats.push_back(t_ns);
pre_stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";";
}
//do inference
tk::dnn::dataDim_t dim = netRT->input_dim;
dim.n = cur_batches;
{
dim.print();
if(TKDNN_VERBOSE) dim.print();
TKDNN_TSTART
netRT->infer(dim, input_d);
TKDNN_TSTOP
dim.print();
if(TKDNN_VERBOSE) dim.print();
stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";";
}
batchDetected.clear();
{
TKDNN_TSTART
postprocess();
for(int bi=0; bi<cur_batches;++bi)
postprocess(bi, mAP);
TKDNN_TSTOP
post_stats.push_back(t_ns);
if(save_times) *times<<t_ns<<"\n";
}
}
}
/**
* Method to draw bounding boxes and labels on a frame.
*
* @param frames original frame to draw bounding box on.
*/
virtual void draw(std::vector<cv::Mat>& frames){};
};
}}