Batch size > 1 for the 3D demo.

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
+37 -14
View File
@@ -1,7 +1,7 @@
#include <iostream> #include <iostream>
#include <signal.h> #include <signal.h>
#include <stdlib.h> /* srand, rand */ #include <stdlib.h> /* srand, rand */
#include <unistd.h> //#include <unistd.h>
#include <mutex> #include <mutex>
#include "CenternetDetection3D.h" #include "CenternetDetection3D.h"
@@ -24,7 +24,12 @@ int main(int argc, char *argv[]) {
std::string net = "dla34_cnet3d_fp32.rt"; std::string net = "dla34_cnet3d_fp32.rt";
if(argc > 1) if(argc > 1)
net = argv[1]; net = argv[1];
#ifdef __linux__
std::string input = "../demo/yolo_test.mp4"; std::string input = "../demo/yolo_test.mp4";
#elif _WIN32
std::string input = "..\\..\\..\\demo\\yolo_test.mp4";
#endif
if(argc > 2) if(argc > 2)
input = argv[2]; input = argv[2];
char ntype = 'c'; char ntype = 'c';
@@ -33,9 +38,18 @@ int main(int argc, char *argv[]) {
int n_classes = 3; int n_classes = 3;
if(argc > 4) if(argc > 4)
n_classes = atoi(argv[4]); n_classes = atoi(argv[4]);
bool show = false; int n_batch = 1;
if(argc > 5) if(argc > 5)
show = atoi(argv[5]); n_batch = atoi(argv[5]);
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");
if(!show) if(!show)
SAVE_RESULT = true; SAVE_RESULT = true;
@@ -57,7 +71,7 @@ int main(int argc, char *argv[]) {
FatalError("Network type not allowed (3rd parameter)\n"); FatalError("Network type not allowed (3rd parameter)\n");
} }
detNN->init(net, n_classes); detNN->init(net, n_classes, n_batch, conf_thresh);
gRun = true; gRun = true;
@@ -75,30 +89,40 @@ int main(int argc, char *argv[]) {
} }
cv::Mat frame; cv::Mat frame;
cv::Mat dnn_input;
if(show) if(show)
cv::namedWindow("detection", cv::WINDOW_NORMAL); cv::namedWindow("detection", cv::WINDOW_NORMAL);
std::vector<tk::dnn::box> detected_bbox; std::vector<cv::Mat> batch_frame;
std::vector<cv::Mat> batch_dnn_input;
while(gRun) { while(gRun) {
batch_dnn_input.clear();
batch_frame.clear();
for(int bi=0; bi< n_batch; ++bi){
cap >> frame; cap >> frame;
if(!frame.data) { if(!frame.data)
break; break;
}
batch_frame.push_back(frame);
// this will be resized to the net format // this will be resized to the net format
dnn_input = frame.clone(); batch_dnn_input.push_back(frame.clone());
}
if(!frame.data)
break;
//inference //inference
detNN->update(dnn_input); detNN->update(batch_dnn_input, n_batch);
frame = detNN->draw(frame); detNN->draw(batch_frame);
if(show){ if(show){
cv::imshow("detection", frame); for(int bi=0; bi< n_batch; ++bi){
cv::imshow("detection", batch_frame[bi]);
cv::waitKey(1); cv::waitKey(1);
} }
if(SAVE_RESULT) }
if(n_batch == 1 && SAVE_RESULT)
resultVideo << frame; resultVideo << frame;
} }
@@ -124,7 +148,6 @@ int main(int argc, char *argv[]) {
std::cout<<"Avg: "<<mean<<" ms\n"<<COL_END; std::cout<<"Avg: "<<mean<<" ms\n"<<COL_END;
return 0; return 0;
} }
+4 -8
View File
@@ -74,22 +74,18 @@ private:
int *ids_out; int *ids_out;
struct threshold op; struct threshold op;
float peakThreshold = 0.2;
float centerThreshold = 0.3; //default 0.5
cv::Mat corners, pts3DHomo; cv::Mat corners, pts3DHomo;
std::vector<box3D> detected3D;
std::vector<int>cls3D;
std::vector<std::vector<int>> face_id; std::vector<std::vector<int>> face_id;
public: public:
CenternetDetection3D() {}; CenternetDetection3D() {};
~CenternetDetection3D() {}; ~CenternetDetection3D() {};
bool init(const std::string& tensor_path, const int n_classes=3); 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); void preprocess(cv::Mat &frame, const int bi=0);
void postprocess(); void postprocess(const int bi=0,const bool mAP=false);
cv::Mat draw(cv::Mat &frame); void draw(std::vector<cv::Mat>& frames);
}; };
+6 -5
View File
@@ -145,6 +145,7 @@ private:
int count_det; int count_det;
//tracks //tracks
std::vector<struct trackingRes> tr_res; std::vector<struct trackingRes> tr_res;
std::vector<std::vector<struct trackingRes>> batchTracked;
int count_tr; int count_tr;
int track_id=0; int track_id=0;
@@ -153,7 +154,7 @@ private:
bool init_pre_inf(); bool init_pre_inf();
bool init_postprocessing(); bool init_postprocessing();
bool init_visualization(const int n_classes); bool init_visualization(const int n_classes);
void pre_inf(); void pre_inf(const int bi);
void _get_additional_inputs(); void _get_additional_inputs();
cv::Mat transform_preds_with_trans(float x1, float x2); cv::Mat transform_preds_with_trans(float x1, float x2);
void tracking(); void tracking();
@@ -162,10 +163,10 @@ public:
tk::dnn::Network *pre_phase_net = nullptr; tk::dnn::Network *pre_phase_net = nullptr;
CenternetDetection3DTrack() {}; CenternetDetection3DTrack() {};
~CenternetDetection3DTrack() {}; ~CenternetDetection3DTrack() {};
bool init(const std::string& tensor_path, const int n_classes=3); 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); void preprocess(cv::Mat &frame, const int bi=0);
void postprocess(); void postprocess(const int bi=0,const bool mAP=false);
cv::Mat draw(cv::Mat &frame); void draw(std::vector<cv::Mat>& frames);
}; };
+48 -27
View File
@@ -4,7 +4,10 @@
#include <iostream> #include <iostream>
#include <signal.h> #include <signal.h>
#include <stdlib.h> #include <stdlib.h>
#ifdef __linux__
#include <unistd.h> #include <unistd.h>
#endif
#include <mutex> #include <mutex>
#include "utils.h" #include "utils.h"
@@ -30,10 +33,12 @@ class DetectionNN3D {
tk::dnn::NetworkRT *netRT = nullptr; tk::dnn::NetworkRT *netRT = nullptr;
dnnType *input_d; dnnType *input_d;
cv::Size originalSize; std::vector<cv::Size> originalSize;
cv::Scalar colors[256]; cv::Scalar colors[256];
int nBatches = 1;
#ifdef OPENCV_CUDACONTRIB #ifdef OPENCV_CUDACONTRIB
cv::cuda::GpuMat bgr[3]; cv::cuda::GpuMat bgr[3];
cv::cuda::GpuMat imagePreproc; cv::cuda::GpuMat imagePreproc;
@@ -47,21 +52,26 @@ class DetectionNN3D {
* This method preprocess the image, before feeding it to the NN. * This method preprocess the image, before feeding it to the NN.
* *
* @param frame original frame to adapt for inference. * @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 * This method postprocess the output of the NN to obtain the correct
* boundig boxes. * 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: public:
int classes = 0; int classes = 0;
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/ 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<double> pre_stats, stats, post_stats, visual_stats; /*keeps track of inference times (ms)*/
std::vector<std::string> classesNames; std::vector<std::string> classesNames;
@@ -69,43 +79,43 @@ class DetectionNN3D {
~DetectionNN3D(){}; ~DetectionNN3D(){};
/** /**
* Method used to inialize the class, allocate memory and compute * Method used to initialize the class, allocate memory and compute
* needed data. * 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_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. * @return true if everything is correct, false otherwise.
*/ */
virtual bool init(const std::string& tensor_path, const int n_classes=3) = 0; 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;
/**
* 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){};
/** /**
* This method performs the whole detection of the NN. * 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 * @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
* box are needed, as in some cases for the mAP calculation.
*/ */
void update(cv::Mat &frame, bool save_times=false, std::ofstream *times=nullptr){ 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(!frame.data)
FatalError("No image data feed to detection");
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");
originalSize = frame.size(); originalSize.clear();
printCenteredTitle(" TENSORRT detection ", '=', 30); if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
{ {
TKDNN_TSTART 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 TKDNN_TSTOP
pre_stats.push_back(t_ns); pre_stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";"; if(save_times) *times<<t_ns<<";";
@@ -113,24 +123,35 @@ class DetectionNN3D {
//do inference //do inference
tk::dnn::dataDim_t dim = netRT->input_dim; tk::dnn::dataDim_t dim = netRT->input_dim;
dim.n = cur_batches;
{ {
dim.print(); if(TKDNN_VERBOSE) dim.print();
TKDNN_TSTART TKDNN_TSTART
netRT->infer(dim, input_d); netRT->infer(dim, input_d);
TKDNN_TSTOP TKDNN_TSTOP
dim.print(); if(TKDNN_VERBOSE) dim.print();
stats.push_back(t_ns); stats.push_back(t_ns);
if(save_times) *times<<t_ns<<";"; if(save_times) *times<<t_ns<<";";
} }
batchDetected.clear();
{ {
TKDNN_TSTART TKDNN_TSTART
postprocess(); for(int bi=0; bi<cur_batches;++bi)
postprocess(bi, mAP);
TKDNN_TSTOP TKDNN_TSTOP
post_stats.push_back(t_ns); post_stats.push_back(t_ns);
if(save_times) *times<<t_ns<<"\n"; 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){};
}; };
}} }}
+31 -28
View File
@@ -3,10 +3,12 @@
namespace tk { namespace dnn { namespace tk { namespace dnn {
bool CenternetDetection3D::init(const std::string& tensor_path, const int n_classes){ bool CenternetDetection3D::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"; std::cout<<(tensor_path).c_str()<<"\n";
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() ); netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
classes = n_classes; classes = n_classes;
nBatches = n_batches;
confThreshold = conf_thresh;
dim = netRT->input_dim; dim = netRT->input_dim;
@@ -28,7 +30,7 @@ bool CenternetDetection3D::init(const std::string& tensor_path, const int n_clas
trans = cv::Mat(cv::Size(3,2), CV_32F); trans = cv::Mat(cv::Size(3,2), CV_32F);
trans2 = cv::Mat(cv::Size(3,2), CV_32F); trans2 = cv::Mat(cv::Size(3,2), CV_32F);
checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot())); checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot() * nBatches));
dim_hm = tk::dnn::dataDim_t(1, 3, 128, 128, 1); dim_hm = tk::dnn::dataDim_t(1, 3, 128, 128, 1);
dim_wh = tk::dnn::dataDim_t(1, 2, 128, 128, 1); dim_wh = tk::dnn::dataDim_t(1, 2, 128, 128, 1);
@@ -91,7 +93,7 @@ bool CenternetDetection3D::init(const std::string& tensor_path, const int n_clas
checkCuda(cudaMemcpy(mean_d, mean, 3*sizeof(float), cudaMemcpyHostToDevice)); checkCuda(cudaMemcpy(mean_d, mean, 3*sizeof(float), cudaMemcpyHostToDevice));
checkCuda(cudaMemcpy(stddev_d, stddev, 3*sizeof(float), cudaMemcpyHostToDevice)); checkCuda(cudaMemcpy(stddev_d, stddev, 3*sizeof(float), cudaMemcpyHostToDevice));
#else #else
checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot())); checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot() * nBatches));
mean << 0.485, 0.456, 0.406; mean << 0.485, 0.456, 0.406;
stddev << 0.229, 0.224, 0.225; stddev << 0.229, 0.224, 0.225;
#endif #endif
@@ -154,13 +156,13 @@ bool CenternetDetection3D::init(const std::string& tensor_path, const int n_clas
// ([[0,1,5,4], [1,2,6, 5], [2,3,7,6], [3,0,4,7]]); // ([[0,1,5,4], [1,2,6, 5], [2,3,7,6], [3,0,4,7]]);
} }
void CenternetDetection3D::preprocess(cv::Mat &frame){ void CenternetDetection3D::preprocess(cv::Mat &frame, const int bi){
// -----------------------------------pre-process ------------------------------------------ // -----------------------------------pre-process ------------------------------------------
// auto start_t = std::chrono::steady_clock::now(); // auto start_t = std::chrono::steady_clock::now();
// auto step_t = std::chrono::steady_clock::now(); // auto step_t = std::chrono::steady_clock::now();
// auto end_t = std::chrono::steady_clock::now(); // auto end_t = std::chrono::steady_clock::now();
cv::Size sz = originalSize; cv::Size sz = originalSize[bi];
// std::cout<<"image: "<<sz.width<<", "<<sz.height<<std::endl; // std::cout<<"image: "<<sz.width<<", "<<sz.height<<std::endl;
cv::Size sz_old; cv::Size sz_old;
float scale = 1.0; float scale = 1.0;
@@ -238,7 +240,7 @@ void CenternetDetection3D::preprocess(cv::Mat &frame){
// std::cout << " TIME normalize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << " us" << std::endl; // std::cout << " TIME normalize: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << " us" << std::endl;
// step_t = end_t; // step_t = end_t;
checkCuda(cudaMemcpy(input_d, d_ptrs, dim2.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice)); checkCuda(cudaMemcpy(input_d+ netRT->input_dim.tot()*bi, d_ptrs, dim2.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice));
// end_t = std::chrono::steady_clock::now(); // end_t = std::chrono::steady_clock::now();
// std::cout << " TIME Memcpy to input_d: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << " us" << std::endl; // std::cout << " TIME Memcpy to input_d: " << std::chrono::duration_cast<std::chrono:: microseconds>(end_t - step_t).count() << " us" << std::endl;
@@ -280,21 +282,21 @@ void CenternetDetection3D::preprocess(cv::Mat &frame){
int idx = i*imageF.rows*imageF.cols; int idx = i*imageF.rows*imageF.cols;
int ch = dim2.c-3 +i; int ch = dim2.c-3 +i;
// std::cout<<"i: "<<i<<", idx: "<<idx<<", ch: "<<ch<<std::endl; // std::cout<<"i: "<<i<<", idx: "<<idx<<", ch: "<<ch<<std::endl;
memcpy((void*)&input[idx], (void*)bgr[ch].data, imageF.rows*imageF.cols*sizeof(dnnType)); memcpy((void*)&input[idx+ netRT->input_dim.tot()*bi], (void*)bgr[ch].data, imageF.rows*imageF.cols*sizeof(dnnType));
} }
checkCuda(cudaMemcpyAsync(input_d, input, dim2.tot()*sizeof(dnnType), cudaMemcpyHostToDevice)); checkCuda(cudaMemcpyAsync(input_d+ netRT->input_dim.tot()*bi, input+ netRT->input_dim.tot()*bi, dim2.tot()*sizeof(dnnType), cudaMemcpyHostToDevice));
#endif #endif
} }
void CenternetDetection3D::postprocess(){ void CenternetDetection3D::postprocess(const int bi, const bool mAP) {
dnnType *rt_out[7]; dnnType *rt_out[7];
rt_out[0] = (dnnType *)netRT->buffersRT[1]; rt_out[0] = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;
rt_out[1] = (dnnType *)netRT->buffersRT[2]; rt_out[1] = (dnnType *)netRT->buffersRT[2]+ netRT->buffersDIM[2].tot()*bi;
rt_out[2] = (dnnType *)netRT->buffersRT[3]; rt_out[2] = (dnnType *)netRT->buffersRT[3]+ netRT->buffersDIM[3].tot()*bi;
rt_out[3] = (dnnType *)netRT->buffersRT[4]; rt_out[3] = (dnnType *)netRT->buffersRT[4]+ netRT->buffersDIM[4].tot()*bi;
rt_out[4] = (dnnType *)netRT->buffersRT[5]; rt_out[4] = (dnnType *)netRT->buffersRT[5]+ netRT->buffersDIM[5].tot()*bi;
rt_out[5] = (dnnType *)netRT->buffersRT[6]; rt_out[5] = (dnnType *)netRT->buffersRT[6]+ netRT->buffersDIM[6].tot()*bi;
rt_out[6] = (dnnType *)netRT->buffersRT[7]; rt_out[6] = (dnnType *)netRT->buffersRT[7]+ netRT->buffersDIM[7].tot()*bi;
// ------------------------------------ process -------------------------------------------- // ------------------------------------ process --------------------------------------------
activationSIGMOIDForward(rt_out[0], rt_out[0], dim_hm.tot()); activationSIGMOIDForward(rt_out[0], rt_out[0], dim_hm.tot());
@@ -404,8 +406,7 @@ void CenternetDetection3D::postprocess(){
if(rot_y<M_PI) if(rot_y<M_PI)
rot_y += 2*M_PI; rot_y += 2*M_PI;
// if(scores[j] > peakThreshold) { if(scores[j] > confThreshold) {
if(scores[j] > centerThreshold) {
if(z>0) { if(z>0) {
// compute_box_3d // compute_box_3d
r.at<float>(0,0) = std::cos(rot_y); r.at<float>(0,0) = std::cos(rot_y);
@@ -457,16 +458,17 @@ void CenternetDetection3D::postprocess(){
} }
res.cl = i; res.cl = i;
res.prob = scores[j]; res.prob = scores[j];
res.print(); //res.print();
detected3D.push_back(res); detected3D.push_back(res);
} }
} }
} }
} }
} }
batchDetected.push_back(detected3D);
} }
cv::Mat CenternetDetection3D::draw(cv::Mat &frame) { void CenternetDetection3D::draw(std::vector<cv::Mat>& frames) {
tk::dnn::box3D b; tk::dnn::box3D b;
int x0, w, x1, y0, h, y1; int x0, w, x1, y0, h, y1;
int objClass; int objClass;
@@ -476,23 +478,24 @@ cv::Mat CenternetDetection3D::draw(cv::Mat &frame) {
float font_scale = 0.5; float font_scale = 0.5;
int thickness = 2; int thickness = 2;
for(int bi=0; bi<frames.size(); ++bi){
// draw dets // draw dets
for(int i=0; i<detected3D.size(); i++) { for(int i=0; i<batchDetected[bi].size(); i++) {
b = detected3D[i]; b = batchDetected[bi][i];
for(int ind_f = 3; ind_f>=0; ind_f--) { for(int ind_f = 3; ind_f>=0; ind_f--) {
for(int j=0; j<4; j++) { for(int j=0; j<4; j++) {
cv::line(frame, cv::Point(b.corners.at(face_id.at(ind_f).at(j) * 2), cv::line(frames[bi], cv::Point(b.corners.at(face_id.at(ind_f).at(j) * 2),
b.corners.at(face_id.at(ind_f).at(j) * 2 + 1)), b.corners.at(face_id.at(ind_f).at(j) * 2 + 1)),
cv::Point(b.corners.at(face_id.at(ind_f).at((j+1)%4) * 2), cv::Point(b.corners.at(face_id.at(ind_f).at((j+1)%4) * 2),
b.corners.at(face_id.at(ind_f).at((j+1)%4) * 2 + 1)), b.corners.at(face_id.at(ind_f).at((j+1)%4) * 2 + 1)),
colors[b.cl], 2); colors[b.cl], 2);
if(ind_f == 0) { if(ind_f == 0) {
cv::line(frame, cv::Point(b.corners.at(face_id.at(ind_f).at(0) * 2), cv::line(frames[bi], cv::Point(b.corners.at(face_id.at(ind_f).at(0) * 2),
b.corners.at(face_id.at(ind_f).at(0) * 2 + 1)), b.corners.at(face_id.at(ind_f).at(0) * 2 + 1)),
cv::Point(b.corners.at(face_id.at(ind_f).at(2) * 2), cv::Point(b.corners.at(face_id.at(ind_f).at(2) * 2),
b.corners.at(face_id.at(ind_f).at(2) * 2 + 1)), colors[b.cl], 2); b.corners.at(face_id.at(ind_f).at(2) * 2 + 1)), colors[b.cl], 2);
cv::line(frame, cv::Point(b.corners.at(face_id.at(ind_f).at(1) * 2), cv::line(frames[bi], cv::Point(b.corners.at(face_id.at(ind_f).at(1) * 2),
b.corners.at(face_id.at(ind_f).at(1) * 2 + 1)), b.corners.at(face_id.at(ind_f).at(1) * 2 + 1)),
cv::Point(b.corners.at(face_id.at(ind_f).at(3) * 2), cv::Point(b.corners.at(face_id.at(ind_f).at(3) * 2),
b.corners.at(face_id.at(ind_f).at(3) * 2 + 1)), colors[b.cl], 2); b.corners.at(face_id.at(ind_f).at(3) * 2 + 1)), colors[b.cl], 2);
@@ -501,15 +504,15 @@ cv::Mat CenternetDetection3D::draw(cv::Mat &frame) {
} }
// draw label // draw label
cv::Size text_size = getTextSize(classesNames[b.cl], cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline); cv::Size text_size = getTextSize(classesNames[b.cl], cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
cv::rectangle(frame, cv::Point(b.corners.at(face_id.at(0).at(0) * 2), cv::rectangle(frames[bi], cv::Point(b.corners.at(face_id.at(0).at(0) * 2),
b.corners.at(face_id.at(0).at(0) * 2 + 1)), b.corners.at(face_id.at(0).at(0) * 2 + 1)),
cv::Point((b.corners.at(face_id.at(0).at(0) * 2) + text_size.width - 2), cv::Point((b.corners.at(face_id.at(0).at(0) * 2) + text_size.width - 2),
(b.corners.at(face_id.at(0).at(0) * 2 + 1)) - text_size.height - 2), colors[b.cl], -1); (b.corners.at(face_id.at(0).at(0) * 2 + 1)) - text_size.height - 2), colors[b.cl], -1);
cv::putText(frame, classesNames[b.cl], cv::Point(b.corners.at(face_id.at(0).at(0) * 2), cv::putText(frames[bi], classesNames[b.cl], cv::Point(b.corners.at(face_id.at(0).at(0) * 2),
b.corners.at(face_id.at(0).at(0) * 2 + 1) - (baseline / 2)), b.corners.at(face_id.at(0).at(0) * 2 + 1) - (baseline / 2)),
cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness); cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), thickness);
} }
return frame; }
} }
}} }}
+87 -84
View File
@@ -3,12 +3,14 @@
namespace tk { namespace dnn { namespace tk { namespace dnn {
bool CenternetDetection3DTrack::init(const std::string& tensor_path, const int n_classes){
std::cout<<(tensor_path).c_str()<<"\n"; bool CenternetDetection3DTrack::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh) {
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() ); netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
dim = netRT->input_dim; dim = netRT->input_dim;
dim.c = 3; dim.c = 3;
nBatches = n_batches;
confThreshold = conf_thresh;
init_preprocessing(); init_preprocessing();
init_pre_inf(); init_pre_inf();
@@ -46,13 +48,13 @@ bool CenternetDetection3DTrack::init_preprocessing(){
checkCuda(cudaMemcpy(mean_d, mean, 3*sizeof(float), cudaMemcpyHostToDevice)); checkCuda(cudaMemcpy(mean_d, mean, 3*sizeof(float), cudaMemcpyHostToDevice));
checkCuda(cudaMemcpy(stddev_d, stddev, 3*sizeof(float), cudaMemcpyHostToDevice)); checkCuda(cudaMemcpy(stddev_d, stddev, 3*sizeof(float), cudaMemcpyHostToDevice));
#else #else
checkCuda(cudaMallocHost(&input, sizeof(dnnType)*dim.tot())); checkCuda(cudaMallocHost(&input, sizeof(dnnType)*dim.tot() * nBatches));
mean << 0.40789655, 0.44719303, 0.47026116; mean << 0.40789655, 0.44719303, 0.47026116;
stddev << 0.2886383, 0.27408165, 0.27809834; stddev << 0.2886383, 0.27408165, 0.27809834;
#endif #endif
checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot())); checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot() * nBatches));
checkCuda(cudaMalloc(&input_pre_inf_d, sizeof(dnnType)*dim.tot())); checkCuda(cudaMalloc(&input_pre_inf_d, sizeof(dnnType)*dim.tot()));
checkCuda( cudaMalloc(&d_ptrs, dim.tot() * sizeof(float)) ); checkCuda( cudaMalloc(&d_ptrs, dim.tot() * sizeof(float)) );
} }
@@ -276,20 +278,20 @@ void CenternetDetection3DTrack::_get_additional_inputs(){
//None no additional input //None no additional input
} }
void CenternetDetection3DTrack::pre_inf(){ void CenternetDetection3DTrack::pre_inf(const int bi){
TKDNN_TSTART TKDNN_TSTART
tk::dnn::dataDim_t dim_aus; tk::dnn::dataDim_t dim_aus;
pre_phase_net->infer(dim_aus, nullptr); pre_phase_net->infer(dim_aus, nullptr);
TKDNN_TSTOP TKDNN_TSTOP
checkCuda( cudaDeviceSynchronize() ); checkCuda( cudaDeviceSynchronize() );
checkCuda( cudaMemcpy(input_d, pre_phase_net->layers[pre_phase_net->num_layers-1]->dstData, netRT->input_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice) ); checkCuda( cudaMemcpy(input_d+ netRT->input_dim.tot()*bi, pre_phase_net->layers[pre_phase_net->num_layers-1]->dstData, netRT->input_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice) );
checkCuda( cudaDeviceSynchronize() ); checkCuda( cudaDeviceSynchronize() );
} }
void CenternetDetection3DTrack::preprocess(cv::Mat &frame){ void CenternetDetection3DTrack::preprocess(cv::Mat &frame, const int bi){
// -----------------------------------pre-process ------------------------------------------ // -----------------------------------pre-process ------------------------------------------
batchTracked.clear();
cv::Size sz = originalSize; cv::Size sz = originalSize[bi];
cv::Size sz_old; cv::Size sz_old;
float scale = 1.0; float scale = 1.0;
float new_height = sz.height * scale; float new_height = sz.height * scale;
@@ -302,7 +304,7 @@ void CenternetDetection3DTrack::preprocess(cv::Mat &frame){
// float s = new_width >= new_height ? new_width : new_height; // float s = new_width >= new_height ? new_width : new_height;
// ----------- get_affine_transform // ----------- get_affine_transform
// rot_rad = pi * 0 / 100 --> 0 // rot_rad = pi * 0 / 100 --> 0
dim.print(); //dim.print();
src.at<float>(0,0)=c[0]; src.at<float>(0,0)=c[0];
src.at<float>(0,1)=c[1]; src.at<float>(0,1)=c[1];
src.at<float>(1,0)=c[0]; src.at<float>(1,0)=c[0];
@@ -389,7 +391,7 @@ void CenternetDetection3DTrack::preprocess(cv::Mat &frame){
checkCuda( cudaDeviceSynchronize() ); checkCuda( cudaDeviceSynchronize() );
iter0=false; iter0=false;
} }
pre_inf(); pre_inf(bi);
checkCuda( cudaMemcpy(img_d, input_pre_inf_d, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice) ); checkCuda( cudaMemcpy(img_d, input_pre_inf_d, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice) );
checkCuda( cudaDeviceSynchronize() ); checkCuda( cudaDeviceSynchronize() );
@@ -587,17 +589,17 @@ void CenternetDetection3DTrack::tracking(){
} }
void CenternetDetection3DTrack::postprocess(){ void CenternetDetection3DTrack::postprocess(const int bi, const bool mAP) {
dnnType *rt_out[9]; dnnType *rt_out[9];
rt_out[0] = (dnnType *)netRT->buffersRT[1]; rt_out[0] = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;
rt_out[1] = (dnnType *)netRT->buffersRT[2]; rt_out[1] = (dnnType *)netRT->buffersRT[2]+ netRT->buffersDIM[2].tot()*bi;
rt_out[2] = (dnnType *)netRT->buffersRT[3]; rt_out[2] = (dnnType *)netRT->buffersRT[3]+ netRT->buffersDIM[3].tot()*bi;
rt_out[3] = (dnnType *)netRT->buffersRT[4]; rt_out[3] = (dnnType *)netRT->buffersRT[4]+ netRT->buffersDIM[4].tot()*bi;
rt_out[4] = (dnnType *)netRT->buffersRT[5]; rt_out[4] = (dnnType *)netRT->buffersRT[5]+ netRT->buffersDIM[5].tot()*bi;
rt_out[5] = (dnnType *)netRT->buffersRT[6]; rt_out[5] = (dnnType *)netRT->buffersRT[6]+ netRT->buffersDIM[6].tot()*bi;
rt_out[6] = (dnnType *)netRT->buffersRT[7]; rt_out[6] = (dnnType *)netRT->buffersRT[7]+ netRT->buffersDIM[7].tot()*bi;
rt_out[7] = (dnnType *)netRT->buffersRT[8]; rt_out[7] = (dnnType *)netRT->buffersRT[8]+ netRT->buffersDIM[8].tot()*bi;
rt_out[8] = (dnnType *)netRT->buffersRT[9]; rt_out[8] = (dnnType *)netRT->buffersRT[9]+ netRT->buffersDIM[9].tot()*bi;
// ------------------------------------ process -------------------------------------------- // ------------------------------------ process --------------------------------------------
@@ -719,78 +721,80 @@ void CenternetDetection3DTrack::postprocess(){
} }
// track step // track step
tracking(); tracking();
batchTracked.push_back(tr_res);
} }
cv::Mat CenternetDetection3DTrack::draw(cv::Mat &frame) { void CenternetDetection3DTrack::draw(std::vector<cv::Mat>& frames) {
struct trackingRes t;
float sc; float sc;
int id; int id;
std::string txt; std::string txt;
int baseline = 0; int baseline = 0;
float font_scale = 0.8; float font_scale = 0.8;
int thickness = 2; int thickness = 2;
for(int i=0; i<count_tr; i++) { for(int bi=0; bi<frames.size(); ++bi) {
id = tr_res[i].tracking_id; // draw dets
txt = classesNames[tr_res[i].det_res.cl-1]+'-'+std::to_string(id); //forse ha bisogno di cl-1 for(int i=0; i<batchTracked[bi].size(); i++) {
t = batchTracked[bi][i];
id = t.tracking_id;
txt = classesNames[t.det_res.cl-1]+'-'+std::to_string(id); //forse ha bisogno di cl-1
cv::Size text_size = getTextSize(txt, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline); cv::Size text_size = getTextSize(txt, cv::FONT_HERSHEY_SIMPLEX, font_scale, thickness, &baseline);
if(tr_res[i].det_res.score > vis_thresh){// && tr_res[i].active!=0) { if(t.det_res.score > vis_thresh){// && t.active!=0) {
if(view2d) { if(view2d) {
cv::rectangle(frames[bi], cv::Point(t.det_res.bb0.at<float>(0,0), t.det_res.bb0.at<float>(0,1)),
cv::Point(t.det_res.bb1.at<float>(0,0), t.det_res.bb1.at<float>(0,1)), tr_colors[t.color], thickness);
cv::rectangle(frames[bi], cv::Point(t.det_res.bb0.at<float>(0,0),
t.det_res.bb0.at<float>(0,1) - text_size.height - thickness),
cv::Point(t.det_res.bb0.at<float>(0,0) + text_size.width,
t.det_res.bb0.at<float>(0,1)), tr_colors[t.color], -1);
cv::putText(frames[bi], txt, cv::Point(t.det_res.bb0.at<float>(0,0),
cv::rectangle(frame, cv::Point(tr_res[i].det_res.bb0.at<float>(0,0), tr_res[i].det_res.bb0.at<float>(0,1)), t.det_res.bb0.at<float>(0,1) - thickness -1),
cv::Point(tr_res[i].det_res.bb1.at<float>(0,0), tr_res[i].det_res.bb1.at<float>(0,1)), tr_colors[tr_res[i].color], thickness);
cv::rectangle(frame, cv::Point(tr_res[i].det_res.bb0.at<float>(0,0),
tr_res[i].det_res.bb0.at<float>(0,1) - text_size.height - thickness),
cv::Point(tr_res[i].det_res.bb0.at<float>(0,0) + text_size.width,
tr_res[i].det_res.bb0.at<float>(0,1)), tr_colors[tr_res[i].color], -1);
cv::putText(frame, txt, cv::Point(tr_res[i].det_res.bb0.at<float>(0,0),
tr_res[i].det_res.bb0.at<float>(0,1) - thickness -1),
cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), 1); cv::FONT_HERSHEY_SIMPLEX, font_scale, cv::Scalar(255, 255, 255), 1);
cv::arrowedLine(frame, cv::Point((int)tr_res[i].det_res.ct.at<float>(0,0), cv::arrowedLine(frames[bi], cv::Point((int)t.det_res.ct.at<float>(0,0),
(int)tr_res[i].det_res.ct.at<float>(0,1)), (int)t.det_res.ct.at<float>(0,1)),
cv::Point((int)(tr_res[i].det_res.ct.at<float>(0,0) + tr_res[i].det_res.tr.at<float>(0,0)), cv::Point((int)(t.det_res.ct.at<float>(0,0) + t.det_res.tr.at<float>(0,0)),
(int)(tr_res[i].det_res.ct.at<float>(0,1) + tr_res[i].det_res.tr.at<float>(0,1))), (int)(t.det_res.ct.at<float>(0,1) + t.det_res.tr.at<float>(0,1))),
cv::Scalar(255, 0, 255), 2); cv::Scalar(255, 0, 255), 2);
} }
//3d //3d
if(!view2d && tr_res[i].det_res.z > 1){ if(!view2d && t.det_res.z > 1){
r.at<float>(0,0) = std::cos(tr_res[i].det_res.rot_y); r.at<float>(0,0) = std::cos(t.det_res.rot_y);
r.at<float>(0,2) = std::sin(tr_res[i].det_res.rot_y); r.at<float>(0,2) = std::sin(t.det_res.rot_y);
r.at<float>(2,0) = -std::sin(tr_res[i].det_res.rot_y); r.at<float>(2,0) = -std::sin(t.det_res.rot_y);
r.at<float>(2,2) = std::cos(tr_res[i].det_res.rot_y); r.at<float>(2,2) = std::cos(t.det_res.rot_y);
corners.at<float>(0,0) = tr_res[i].det_res.dim[2]/2; corners.at<float>(0,0) = t.det_res.dim[2]/2;
corners.at<float>(0,1) = tr_res[i].det_res.dim[2]/2; corners.at<float>(0,1) = t.det_res.dim[2]/2;
corners.at<float>(0,2) = -tr_res[i].det_res.dim[2]/2; corners.at<float>(0,2) = -t.det_res.dim[2]/2;
corners.at<float>(0,3) = -tr_res[i].det_res.dim[2]/2; corners.at<float>(0,3) = -t.det_res.dim[2]/2;
corners.at<float>(0,4) = tr_res[i].det_res.dim[2]/2; corners.at<float>(0,4) = t.det_res.dim[2]/2;
corners.at<float>(0,5) = tr_res[i].det_res.dim[2]/2; corners.at<float>(0,5) = t.det_res.dim[2]/2;
corners.at<float>(0,6) = -tr_res[i].det_res.dim[2]/2; corners.at<float>(0,6) = -t.det_res.dim[2]/2;
corners.at<float>(0,7) = -tr_res[i].det_res.dim[2]/2; corners.at<float>(0,7) = -t.det_res.dim[2]/2;
corners.at<float>(1,4) = -tr_res[i].det_res.dim[0]; corners.at<float>(1,4) = -t.det_res.dim[0];
corners.at<float>(1,5) = -tr_res[i].det_res.dim[0]; corners.at<float>(1,5) = -t.det_res.dim[0];
corners.at<float>(1,6) = -tr_res[i].det_res.dim[0]; corners.at<float>(1,6) = -t.det_res.dim[0];
corners.at<float>(1,7) = -tr_res[i].det_res.dim[0]; corners.at<float>(1,7) = -t.det_res.dim[0];
corners.at<float>(2,0) = tr_res[i].det_res.dim[1]/2; corners.at<float>(2,0) = t.det_res.dim[1]/2;
corners.at<float>(2,1) = -tr_res[i].det_res.dim[1]/2; corners.at<float>(2,1) = -t.det_res.dim[1]/2;
corners.at<float>(2,2) = -tr_res[i].det_res.dim[1]/2; corners.at<float>(2,2) = -t.det_res.dim[1]/2;
corners.at<float>(2,3) = tr_res[i].det_res.dim[1]/2; corners.at<float>(2,3) = t.det_res.dim[1]/2;
corners.at<float>(2,4) = tr_res[i].det_res.dim[1]/2; corners.at<float>(2,4) = t.det_res.dim[1]/2;
corners.at<float>(2,5) = -tr_res[i].det_res.dim[1]/2; corners.at<float>(2,5) = -t.det_res.dim[1]/2;
corners.at<float>(2,6) = -tr_res[i].det_res.dim[1]/2; corners.at<float>(2,6) = -t.det_res.dim[1]/2;
corners.at<float>(2,7) = tr_res[i].det_res.dim[1]/2; corners.at<float>(2,7) = t.det_res.dim[1]/2;
cv::Mat aus = r * corners; cv::Mat aus = r * corners;
for(int k=0; k<8; k++) { for(int k=0; k<8; k++) {
aus.at<float>(0,k) += tr_res[i].det_res.x; aus.at<float>(0,k) += t.det_res.x;
aus.at<float>(1,k) += tr_res[i].det_res.y; aus.at<float>(1,k) += t.det_res.y;
aus.at<float>(2,k) += tr_res[i].det_res.z; aus.at<float>(2,k) += t.det_res.z;
} }
// corners.copyTo(pts3DHomo(cv::Rect(0, 0, 8, 3))); // corners.copyTo(pts3DHomo(cv::Rect(0, 0, 8, 3)));
@@ -809,20 +813,20 @@ cv::Mat CenternetDetection3DTrack::draw(cv::Mat &frame) {
aus.release(); aus.release();
for(int ind_f = 3; ind_f>=0; ind_f--) { for(int ind_f = 3; ind_f>=0; ind_f--) {
for(int j=0; j<4; j++) { for(int j=0; j<4; j++) {
cv::line(frame, cv::Point((int)res_corners.at(face_id.at(ind_f).at(j) * 2), cv::line(frames[bi], cv::Point((int)res_corners.at(face_id.at(ind_f).at(j) * 2),
(int)res_corners.at(face_id.at(ind_f).at(j) * 2 + 1)), (int)res_corners.at(face_id.at(ind_f).at(j) * 2 + 1)),
cv::Point((int)res_corners.at(face_id.at(ind_f).at((j+1)%4) * 2), cv::Point((int)res_corners.at(face_id.at(ind_f).at((j+1)%4) * 2),
(int)res_corners.at(face_id.at(ind_f).at((j+1)%4) * 2 + 1)), (int)res_corners.at(face_id.at(ind_f).at((j+1)%4) * 2 + 1)),
tr_colors[tr_res[i].color], 2); tr_colors[t.color], 2);
if(ind_f == 0 && j==3) { if(ind_f == 0 && j==3) {
cv::line(frame, cv::Point((int)res_corners.at(face_id.at(ind_f).at(0) * 2), cv::line(frames[bi], cv::Point((int)res_corners.at(face_id.at(ind_f).at(0) * 2),
(int)res_corners.at(face_id.at(ind_f).at(0) * 2 + 1)), (int)res_corners.at(face_id.at(ind_f).at(0) * 2 + 1)),
cv::Point((int)res_corners.at(face_id.at(ind_f).at(2) * 2), cv::Point((int)res_corners.at(face_id.at(ind_f).at(2) * 2),
(int)res_corners.at(face_id.at(ind_f).at(2) * 2 + 1)), tr_colors[tr_res[i].color], 2); (int)res_corners.at(face_id.at(ind_f).at(2) * 2 + 1)), tr_colors[t.color], 2);
cv::line(frame, cv::Point((int)res_corners.at(face_id.at(ind_f).at(1) * 2), cv::line(frames[bi], cv::Point((int)res_corners.at(face_id.at(ind_f).at(1) * 2),
(int)res_corners.at(face_id.at(ind_f).at(1) * 2 + 1)), (int)res_corners.at(face_id.at(ind_f).at(1) * 2 + 1)),
cv::Point((int)res_corners.at(face_id.at(ind_f).at(3) * 2), cv::Point((int)res_corners.at(face_id.at(ind_f).at(3) * 2),
(int)res_corners.at(face_id.at(ind_f).at(3) * 2 + 1)), tr_colors[tr_res[i].color], 2); (int)res_corners.at(face_id.at(ind_f).at(3) * 2 + 1)), tr_colors[t.color], 2);
} }
} }
} }
@@ -840,22 +844,21 @@ cv::Mat CenternetDetection3DTrack::draw(cv::Mat &frame) {
} }
// if(not no_bbox): // if(not no_bbox):
// cv::rectangle(frame, cv::Point(bb0, bb2), cv::Point(bb1, bb3), // cv::rectangle(frame, cv::Point(bb0, bb2), cv::Point(bb1, bb3),
// tr_colors[tr_res[i].color], thickness); // tr_colors[t.color], thickness);
cv::rectangle(frame, cv::Point(bb0, bb2 - text_size.height - thickness), cv::rectangle(frames[bi], cv::Point(bb0, bb2 - text_size.height - thickness),
cv::Point(bb0 + text_size.width, bb2), tr_colors[tr_res[i].color], -1); cv::Point(bb0 + text_size.width, bb2), tr_colors[t.color], -1);
cv::putText(frame, txt, cv::Point(bb0, bb2 - thickness -1), cv::FONT_HERSHEY_SIMPLEX, cv::putText(frames[bi], txt, cv::Point(bb0, bb2 - thickness -1), cv::FONT_HERSHEY_SIMPLEX,
font_scale, cv::Scalar(255, 255, 255), 1); font_scale, cv::Scalar(255, 255, 255), 1);
cv::arrowedLine(frame, cv::Point((int)((bb0 + bb1)/2), (int)((bb2 + bb3)/2)), cv::arrowedLine(frames[bi], cv::Point((int)((bb0 + bb1)/2), (int)((bb2 + bb3)/2)),
cv::Point((int)((bb0 + bb1)/2 + tr_res[i].det_res.tr.at<float>(0,0)), cv::Point((int)((bb0 + bb1)/2 + t.det_res.tr.at<float>(0,0)),
(int)((bb2 + bb3)/2 + tr_res[i].det_res.tr.at<float>(0,1))), (int)((bb2 + bb3)/2 + t.det_res.tr.at<float>(0,1))),
cv::Scalar(255, 0, 255), 2); cv::Scalar(255, 0, 255), 2);
} }
} }
} }
return frame; }
} }
}} }}