Skip to content

Commit ba8c282

Browse files
sapienzadavidesapienzadavide
committed
Add 3D CenterNet detection class
Signed-off-by: Davide Sapienza <[email protected]>
1 parent 64098ad commit ba8c282

File tree

6 files changed

+779
-0
lines changed

6 files changed

+779
-0
lines changed

‎include/tkDNN/CenternetDetection3D.h‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
#ifndef CENTERNETDETECTION3D_H
2+
#define CENTERNETDETECTION3D_H
3+
4+
#include "kernels.h"
5+
#include <opencv2/videoio.hpp>
6+
#include "opencv2/opencv.hpp"
7+
#include <time.h>
8+
#include <vector>
9+
#include <numeric> // std::iota
10+
#include <algorithm> // std::sort
11+
12+
#include "DetectionNN3D.h"
13+
14+
#include "kernelsThrust.h"
15+
16+
17+
namespace tk { namespace dnn {
18+
19+
class CenternetDetection3D : public DetectionNN3D
20+
{
21+
private:
22+
tk::dnn::dataDim_t dim;
23+
tk::dnn::dataDim_t dim2;
24+
tk::dnn::dataDim_t dim_hm;
25+
tk::dnn::dataDim_t dim_wh;
26+
tk::dnn::dataDim_t dim_reg;
27+
tk::dnn::dataDim_t dim_dep;
28+
tk::dnn::dataDim_t dim_rot;
29+
tk::dnn::dataDim_t dim_dim;
30+
float *topk_scores;
31+
int *topk_inds_;
32+
float *topk_ys_;
33+
float *topk_xs_;
34+
int *ids_d, *ids_;
35+
36+
float *ones;
37+
38+
float *scores, *scores_d;
39+
int *clses, *clses_d;
40+
int *topk_inds_d;
41+
float *topk_ys_d;
42+
float *topk_xs_d;
43+
int *inttopk_xs_d, *inttopk_ys_d;
44+
45+
float *xs, *ys;
46+
47+
float *dep, *rot, *dim_, *wh;
48+
float *dep_d, *rot_d, *dim_d, *wh_d;
49+
50+
float *target_coords;
51+
52+
#ifdef OPENCV_CUDACONTRIB
53+
float *mean_d;
54+
float *stddev_d;
55+
#else
56+
cv::Vec<float, 3> mean;
57+
cv::Vec<float, 3> stddev;
58+
dnnType *input;
59+
#endif
60+
cv::Mat r;
61+
cv::Mat calibs;
62+
float *d_ptrs;
63+
64+
cv::Mat src;
65+
cv::Mat dst;
66+
cv::Mat dst2;
67+
cv::Mat trans, trans2;
68+
//processing
69+
int K = 100;
70+
int width = 128;//56; // TODO
71+
72+
// pointer used in the kernels
73+
float *src_out;
74+
int *ids_out;
75+
76+
struct threshold op;
77+
float peakThreshold = 0.2;
78+
float centerThreshold = 0.3; //default 0.5
79+
cv::Mat corners, pts3DHomo;
80+
81+
std::vector<box3D> detected3D;
82+
std::vector<int>cls3D;
83+
std::vector<std::vector<int>> face_id;
84+
85+
public:
86+
CenternetDetection3D() {};
87+
~CenternetDetection3D() {};
88+
89+
bool init(const std::string& tensor_path, const int n_classes=3);
90+
void preprocess(cv::Mat &frame);
91+
void postprocess();
92+
cv::Mat draw(cv::Mat &frame);
93+
};
94+
95+
96+
} // namespace dnn
97+
} // namespace tk
98+
99+
100+
#endif /*CENTERNETDETECTION_H*/

‎include/tkDNN/DetectionNN3D.h‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
#ifndef DETECTIONNN3D_H
2+
#define DETECTIONNN3D_H
3+
4+
#include <iostream>
5+
#include <signal.h>
6+
#include <stdlib.h>
7+
#include <unistd.h>
8+
#include <mutex>
9+
#include "utils.h"
10+
11+
#include <opencv2/core/core.hpp>
12+
#include <opencv2/highgui/highgui.hpp>
13+
#include <opencv2/imgproc/imgproc.hpp>
14+
15+
#include "tkdnn.h"
16+
17+
// #define OPENCV_CUDACONTRIB //if OPENCV has been compiled with CUDA and contrib.
18+
19+
#ifdef OPENCV_CUDACONTRIB
20+
#include <opencv2/cudawarping.hpp>
21+
#include <opencv2/cudaarithm.hpp>
22+
#endif
23+
24+
25+
namespace tk { namespace dnn {
26+
27+
class DetectionNN3D {
28+
29+
protected:
30+
tk::dnn::NetworkRT *netRT = nullptr;
31+
dnnType *input_d;
32+
33+
cv::Size originalSize;
34+
35+
cv::Scalar colors[256];
36+
37+
#ifdef OPENCV_CUDACONTRIB
38+
cv::cuda::GpuMat bgr[3];
39+
cv::cuda::GpuMat imagePreproc;
40+
#else
41+
cv::Mat bgr[3];
42+
cv::Mat imagePreproc;
43+
dnnType *input;
44+
#endif
45+
46+
/**
47+
* This method preprocess the image, before feeding it to the NN.
48+
*
49+
* @param frame original frame to adapt for inference.
50+
*/
51+
virtual void preprocess(cv::Mat &frame) = 0;
52+
53+
/**
54+
* This method postprocess the output of the NN to obtain the correct
55+
* boundig boxes.
56+
*
57+
*/
58+
virtual void postprocess() = 0;
59+
60+
public:
61+
int classes = 0;
62+
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
63+
64+
std::vector<tk::dnn::box> detected; /*bounding boxes in output*/
65+
std::vector<double> stats; /*keeps track of inference times (ms)*/
66+
std::vector<std::string> classesNames;
67+
68+
DetectionNN3D() {};
69+
~DetectionNN3D(){};
70+
71+
/**
72+
* Method used to inialize the class, allocate memory and compute
73+
* needed data.
74+
*
75+
* @param tensor_path path to the rt file og the NN.
76+
* @param n_classes number of classes for the given dataset.
77+
* @return true if everything is correct, false otherwise.
78+
*/
79+
virtual bool init(const std::string& tensor_path, const int n_classes=3) = 0;
80+
81+
/**
82+
* Method to draw boundixg boxes and labels on a frame.
83+
*
84+
* @param frame orginal frame to draw bounding box on.
85+
* @return frame with boundig boxes.
86+
*/
87+
virtual cv::Mat draw(cv::Mat &frame){};
88+
89+
/**
90+
* This method performs the whole detection of the NN.
91+
*
92+
* @param frame frame to run detection on.
93+
* @param save_times if set to true, preprocess, inference and postprocess times
94+
* are saved on a csv file, otherwise not.
95+
* @param times pointer to the output stream where to write times
96+
*/
97+
void update(cv::Mat &frame, bool save_times=false, std::ofstream *times=nullptr){
98+
if(!frame.data)
99+
FatalError("No image data feed to detection");
100+
101+
if(save_times && times==nullptr)
102+
FatalError("save_times set to true, but no valid ofstream given");
103+
104+
originalSize = frame.size();
105+
printCenteredTitle(" TENSORRT detection ", '=', 30);
106+
{
107+
TIMER_START
108+
preprocess(frame);
109+
TIMER_STOP
110+
if(save_times) *times<<t_ns<<";";
111+
}
112+
113+
//do inference
114+
tk::dnn::dataDim_t dim = netRT->input_dim;
115+
{
116+
dim.print();
117+
TIMER_START
118+
netRT->infer(dim, input_d);
119+
TIMER_STOP
120+
dim.print();
121+
stats.push_back(t_ns);
122+
if(save_times) *times<<t_ns<<";";
123+
}
124+
125+
{
126+
TIMER_START
127+
postprocess();
128+
TIMER_STOP
129+
if(save_times) *times<<t_ns<<"\n";
130+
}
131+
}
132+
};
133+
134+
}}
135+
136+
#endif /* DETECTIONNN3D_H*/

‎include/tkDNN/Layer.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -546,6 +546,16 @@ struct sortable_bbox {
546546
int cl;
547547
float **probs;
548548
};
549+
struct box3D {
550+
int cl;
551+
std::vector<float> corners;
552+
float prob;
553+
554+
void print()
555+
{
556+
std::cout<<"\tcl: "<<cl<<"\tprob: "<<prob<<"\tshape corners: "<<corners.size()<<std::endl;
557+
}
558+
};
549559

550560
/**
551561
Yolo3 layer

‎include/tkDNN/kernelsThrust.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,13 @@ void topk(dnnType *src_begin, int *idsrc, int K, float *topk_scores,
2929
int *topk_inds, float *topk_ys, float *topk_xs);
3030
// void sortAndTopKonDevice(dnnType *src_begin, int *idsrc, float *topk_scores, int *topk_inds, float *topk_ys, float *topk_xs, const int size, const int K, const int n_classes);
3131
void normalize(float *bgr, const int ch, const int h, const int w, const float *mean, const float *stddev);
32+
void transformDep(float *src_begin, float *src_end, float *dst_begin, float *dst_end);
3233
void subtractWithThreshold(dnnType *src_begin, dnnType *src_end, dnnType *src2_begin, dnnType *src_out, struct threshold op);
3334
void topKxyclasses(int *ids_begin, int *ids_end, const int K, const int size, const int wh, int *clses, int *xs, int *ys);
3435
void topKxyAddOffset(int * ids_begin, const int K, const int size, int *intxs_begin, int *intys_begin,
3536
float *xs_begin, float *ys_begin, dnnType *src_begin, float *src_out, int *ids_out);
3637
void bboxes(int * ids_begin, const int K, const int size, float *xs_begin, float *ys_begin,
3738
dnnType *src_begin, float *bbx0, float *bbx1, float *bby0, float *bby1, float *src_out, int *ids_out);
39+
void getRecordsFromTopKId(int * ids_begin, const int K, const int ch, const int size, dnnType *src_begin, float *src_out, int *ids_out);
3840

3941
#endif //KERNELSTHRUST_H

0 commit comments

Comments
 (0)