Add 3D CenterNet detection class

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
2020-05-27 18:03:09 +02:00
parent 64098ad244
commit ba8c282384
6 changed files with 779 additions and 0 deletions
@@ -0,0 +1,100 @@
+#ifndef CENTERNETDETECTION3D_H
+#define CENTERNETDETECTION3D_H
+
+#include "kernels.h"
+#include <opencv2/videoio.hpp>
+#include "opencv2/opencv.hpp"
+#include <time.h>
+#include <vector>
+#include <numeric>      // std::iota
+#include <algorithm>    // std::sort
+
+#include "DetectionNN3D.h"
+
+#include "kernelsThrust.h"
+
+
+namespace tk { namespace dnn { 
+
+class CenternetDetection3D : public DetectionNN3D
+{
+private:
+    tk::dnn::dataDim_t dim;
+    tk::dnn::dataDim_t dim2;
+    tk::dnn::dataDim_t dim_hm; 
+    tk::dnn::dataDim_t dim_wh; 
+    tk::dnn::dataDim_t dim_reg;
+    tk::dnn::dataDim_t dim_dep;
+    tk::dnn::dataDim_t dim_rot;
+    tk::dnn::dataDim_t dim_dim;
+    float *topk_scores;
+    int *topk_inds_;
+    float *topk_ys_;
+    float *topk_xs_;
+    int *ids_d, *ids_;
+
+    float *ones;
+
+    float *scores, *scores_d;
+    int *clses, *clses_d; 
+    int *topk_inds_d;
+    float *topk_ys_d; 
+    float *topk_xs_d;
+    int *inttopk_xs_d, *inttopk_ys_d;
+
+    float *xs, *ys;
+
+    float *dep, *rot, *dim_, *wh;
+    float *dep_d, *rot_d, *dim_d, *wh_d;
+    
+    float *target_coords;
+
+    #ifdef OPENCV_CUDACONTRIB
+        float *mean_d;
+        float *stddev_d;
+    #else
+        cv::Vec<float, 3> mean;
+        cv::Vec<float, 3> stddev;
+        dnnType *input;
+    #endif
+    cv::Mat r;
+    cv::Mat calibs;
+    float *d_ptrs;
+
+    cv::Mat src;
+    cv::Mat dst;
+    cv::Mat dst2;  
+    cv::Mat trans, trans2;
+    //processing
+    int K = 100;
+    int width = 128;//56;        // TODO
+
+    // pointer used in the kernels
+    float *src_out;
+    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);
+};
+
+
+} // namespace dnn
+} // namespace tk
+
+
+#endif /*CENTERNETDETECTION_H*/
@@ -0,0 +1,136 @@
+#ifndef DETECTIONNN3D_H
+#define DETECTIONNN3D_H
+
+#include <iostream>
+#include <signal.h>
+#include <stdlib.h>    
+#include <unistd.h>
+#include <mutex>
+#include "utils.h"
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include "tkdnn.h"
+
+// #define OPENCV_CUDACONTRIB //if OPENCV has been compiled with CUDA and contrib.
+
+#ifdef OPENCV_CUDACONTRIB
+#include <opencv2/cudawarping.hpp>
+#include <opencv2/cudaarithm.hpp>
+#endif
+
+
+namespace tk { namespace dnn {
+
+class DetectionNN3D {
+
+    protected:
+        tk::dnn::NetworkRT *netRT = nullptr;
+        dnnType *input_d;
+
+        cv::Size originalSize;
+
+        cv::Scalar colors[256];
+
+#ifdef OPENCV_CUDACONTRIB
+        cv::cuda::GpuMat bgr[3];
+        cv::cuda::GpuMat imagePreproc;
+#else
+        cv::Mat bgr[3];
+        cv::Mat imagePreproc;
+        dnnType *input;
+#endif
+
+        /**
+         * This method preprocess the image, before feeding it to the NN.
+         *
+         * @param frame original frame to adapt for inference.
+         */
+        virtual void preprocess(cv::Mat &frame) = 0;
+
+        /**
+         * This method postprocess the output of the NN to obtain the correct 
+         * boundig boxes. 
+         * 
+         */
+        virtual void postprocess() = 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<double> stats; /*keeps track of inference times (ms)*/
+        std::vector<std::string> classesNames;
+
+        DetectionNN3D() {};
+        ~DetectionNN3D(){};
+
+        /**
+         * Method used to inialize the class, allocate memory and compute 
+         * needed data.
+         * 
+         * @param tensor_path path to the rt file og the NN.
+         * @param n_classes number of classes for the given dataset.
+         * @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){};
+
+        /**
+         * This method performs the whole detection of the NN.
+         * 
+         * @param frame frame to run detection on.
+         * @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
+         */
+        void update(cv::Mat &frame, bool save_times=false, std::ofstream *times=nullptr){
+            if(!frame.data)
+                FatalError("No image data feed to detection");
+
+            if(save_times && times==nullptr)
+                FatalError("save_times set to true, but no valid ofstream given");
+
+            originalSize = frame.size();
+            printCenteredTitle(" TENSORRT detection ", '=', 30); 
+            {
+                TIMER_START
+                preprocess(frame);
+                TIMER_STOP
+                if(save_times) *times<<t_ns<<";";
+            }
+
+            //do inference
+            tk::dnn::dataDim_t dim = netRT->input_dim;
+            {
+                dim.print();
+                TIMER_START
+                netRT->infer(dim, input_d);
+                TIMER_STOP
+                dim.print();
+                stats.push_back(t_ns);
+                if(save_times) *times<<t_ns<<";";
+            }
+
+            {
+                TIMER_START
+                postprocess();
+                TIMER_STOP
+                if(save_times) *times<<t_ns<<"\n";
+            }
+        }              
+};
+
+}}
+
+#endif /* DETECTIONNN3D_H*/
@@ -546,6 +546,16 @@ struct sortable_bbox {
    int cl;
    float **probs;
 };
+struct box3D {
+    int cl;
+    std::vector<float> corners;
+    float prob;
+
+    void print() 
+    {
+        std::cout<<"\tcl: "<<cl<<"\tprob: "<<prob<<"\tshape corners: "<<corners.size()<<std::endl;
+    }
+};

 /**
    Yolo3 layer
@@ -29,11 +29,13 @@ void topk(dnnType *src_begin, int *idsrc, int K, float *topk_scores,
            int *topk_inds, float *topk_ys, float *topk_xs);
 // 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);            
 void normalize(float *bgr, const int ch, const int h, const int w, const float *mean, const float *stddev);
+void transformDep(float *src_begin, float *src_end, float *dst_begin, float *dst_end);
 void subtractWithThreshold(dnnType *src_begin, dnnType *src_end, dnnType *src2_begin, dnnType *src_out, struct threshold op);
 void topKxyclasses(int *ids_begin, int *ids_end, const int K, const int size, const int wh, int *clses, int *xs, int *ys);
 void topKxyAddOffset(int * ids_begin, const int K, const int size, int *intxs_begin, int *intys_begin, 
                     float *xs_begin, float *ys_begin, dnnType *src_begin, float *src_out, int *ids_out);
 void bboxes(int * ids_begin, const int K, const int size, float *xs_begin, float *ys_begin, 
            dnnType *src_begin, float *bbx0, float *bbx1, float *bby0, float *bby1, float *src_out, int *ids_out);
+void getRecordsFromTopKId(int * ids_begin, const int K, const int ch, const int size, dnnType *src_begin, float *src_out, int *ids_out);

 #endif //KERNELSTHRUST_H