Add the calibration matrix reading for CenterTrack

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>

demo/demo/demo3D.cpp

@@ -70,8 +70,17 @@ int main(int argc, char *argv[]) {
         default:
         FatalError("Network type not allowed (3rd parameter)\n");
     }
-
+    std::vector<cv::Mat> calibs;
-    detNN->init(net, n_classes, n_batch, conf_thresh);
+    // cv::Mat calib = cv::Mat::zeros(cv::Size(3,3), CV_32F);
+    // calib.at<float>(0,0) = 864.1243196486207;// * 512.0;//884.081444212;//864.1243196486207 * 512.0;// 633.0;
+    // calib.at<float>(0,2) = 726.7271690557819;// * 512.0;//0.0;//726.7271690557819 * 512.0;// 0.0; //w/2
+    // calib.at<float>(1,1) = 883.6552349216504;// * 512.0;//884.081444212;//883.6552349216504 * 512.0;// 633.0;
+    // calib.at<float>(1,2) = 506.8548506986564;// * 512.0;//0.0;//506.8548506986564 * 512.0;// 0.0; //h/2
+    // calibs.push_back(calib);
+    // calibs.push_back(calib);
+    // calibs.push_back(calib);
+    // calibs.push_back(calib);
+    detNN->init(net, n_classes, n_batch, conf_thresh, calibs);
 
     gRun = true;
 
@@ -87,7 +96,8 @@ int main(int argc, char *argv[]) {
         int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
         resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M','P','4','V'), 30, cv::Size(w, h));
     }
-
+    cv::Size sz_resize = cv::Size(512,512); 
+    std::vector<cv::Size> sz_orig;
     cv::Mat frame;
     if(show)
 	    cv::namedWindow("detection", cv::WINDOW_NORMAL);
@@ -98,12 +108,15 @@ int main(int argc, char *argv[]) {
     while(gRun) {
         batch_dnn_input.clear();
         batch_frame.clear();
+        sz_orig.clear();
         
         for(int bi=0; bi< n_batch; ++bi){
             cap >> frame; 
             if(!frame.data) 
                 break;
-            
+            sz_orig.push_back(frame.size());
+            if(calibs.size() != 0)
+                resize(frame, frame, sz_resize);
             batch_frame.push_back(frame);
 
             // this will be resized to the net format
@@ -113,11 +126,13 @@ int main(int argc, char *argv[]) {
             break;  
  
         //inference
-        detNN->update(batch_dnn_input, n_batch);
+        detNN->update(batch_dnn_input, n_batch, false, nullptr, false, sz_orig);
         detNN->draw(batch_frame);
 
         if(show){
             for(int bi=0; bi< n_batch; ++bi){
+                if(calibs.size() != 0)
+                    resize(batch_frame[bi], batch_frame[bi], sz_orig[bi]);
                 cv::imshow("detection", batch_frame[bi]);
                 cv::waitKey(1);
             }

include/tkDNN/CenternetDetection3D.h

@@ -82,8 +82,8 @@ public:
     CenternetDetection3D() {};
     ~CenternetDetection3D() {}; 
 
-    bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1, const float conf_thresh=0.3);
+    bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1, const float conf_thresh=0.3, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>());
-    void preprocess(cv::Mat &frame, const int bi=0);
+    void preprocess(cv::Mat &frame, const int bi=0, const std::vector<cv::Size>& stream_size=std::vector<cv::Size>());
     void postprocess(const int bi=0,const bool mAP=false);
     void draw(std::vector<cv::Mat>& frames);
 };

include/tkDNN/CenternetDetection3DTrack.h

@@ -74,6 +74,10 @@ private:
     #endif
     float *d_ptrs;
 
+    std::vector<cv::Mat> inputCalibs;
+
+    std::vector<cv::Size> sz_old;
+
     cv::Mat src;
     cv::Mat dst;
     cv::Mat dst2;  
@@ -124,7 +128,7 @@ private:
 
     /* visualization */
     cv::Mat r;
-    cv::Mat calibs;
+    std::vector<cv::Mat> calibs;
     cv::Mat corners, pts3DHomo;  
 
     std::vector<std::vector<int>> face_id;
@@ -163,8 +167,8 @@ public:
     tk::dnn::Network *pre_phase_net = nullptr;
     CenternetDetection3DTrack() {};
     ~CenternetDetection3DTrack() {};
-    bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1, const float conf_thresh=0.3);
+    bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1, const float conf_thresh=0.3, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>());
-    void preprocess(cv::Mat &frame, const int bi=0);
+    void preprocess(cv::Mat &frame, const int bi=0, const std::vector<cv::Size>& stream_size=std::vector<cv::Size>());
     void postprocess(const int bi=0,const bool mAP=false);
     void draw(std::vector<cv::Mat>& frames);
 };

include/tkDNN/DetectionNN3D.h

@@ -54,7 +54,7 @@ class DetectionNN3D {
          * @param frame original frame to adapt for inference.
          * @param bi batch index
          */
-        virtual void preprocess(cv::Mat &frame, const int bi=0) = 0;
+        virtual void preprocess(cv::Mat &frame, const int bi=0 , const std::vector<cv::Size>& stream_size=std::vector<cv::Size>()) = 0;
 
         /**
          * This method postprocess the output of the NN to obtain the correct 
@@ -87,7 +87,8 @@ class DetectionNN3D {
          * @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, const int n_batches=1, const float conf_thresh=0.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, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>()) = 0;
 
         /**
          * This method performs the whole detection of the NN.
@@ -100,7 +101,8 @@ class DetectionNN3D {
          * @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(std::vector<cv::Mat>& frames, const int cur_batches=1, 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, const std::vector<cv::Size>& stream_size=std::vector<cv::Size>()){
             if(save_times && times==nullptr)
                 FatalError("save_times set to true, but no valid ofstream given");
             if(cur_batches > nBatches)
@@ -114,7 +116,7 @@ class DetectionNN3D {
                     if(!frames[bi].data)
                         FatalError("No image data feed to detection");
                     originalSize.push_back(frames[bi].size());
-                    preprocess(frames[bi], bi);   
+                    preprocess(frames[bi], bi, stream_size);   
                 }
                 TKDNN_TSTOP
                 pre_stats.push_back(t_ns); 

src/CenternetDetection3D.cpp

@@ -3,7 +3,8 @@
 
 namespace tk { namespace dnn {
 
-bool CenternetDetection3D::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh) {
+bool CenternetDetection3D::init(const std::string& tensor_path, const int n_classes, const int n_batches, 
+                                const float conf_thresh, const std::vector<cv::Mat>& k_calibs) {
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
     classes = n_classes;
@@ -156,7 +157,7 @@ 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]]);
 }
 
-void CenternetDetection3D::preprocess(cv::Mat &frame, const int bi){
+void CenternetDetection3D::preprocess(cv::Mat &frame, const int bi, const std::vector<cv::Size>& stream_size){
      // -----------------------------------pre-process ------------------------------------------
     
     // auto start_t = std::chrono::steady_clock::now();

src/CenternetDetection3DTrack.cpp

@@ -4,14 +4,15 @@
 namespace tk { namespace dnn {
 
 
-bool CenternetDetection3DTrack::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh) {
+bool CenternetDetection3DTrack::init(const std::string& tensor_path, const int n_classes, const int n_batches, 
+                                    const float conf_thresh, const std::vector<cv::Mat>& k_calibs) {
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
     
     dim = netRT->input_dim;
     dim.c = 3;
     nBatches = n_batches;
     confThreshold = conf_thresh;
-
+    inputCalibs = k_calibs;
     init_preprocessing();
     init_pre_inf();
     init_postprocessing();
@@ -37,7 +38,10 @@ bool CenternetDetection3DTrack::init_preprocessing(){
     dst2.at<float>(2,0)=dst2.at<float>(1,0) + (-dst2.at<float>(0,1)+dst2.at<float>(1,1) );
     dst2.at<float>(2,1)=dst2.at<float>(1,1) + (dst2.at<float>(0,0)-dst2.at<float>(1,0) );
 
-
+    for(int bi=0; bi<nBatches; bi++) {
+        sz_old.push_back(cv::Size(0,0));
+    }
+    
 #ifdef OPENCV_CUDACONTRIB
 
     checkCuda( cudaMalloc(&mean_d, 3 * sizeof(float)) );
@@ -183,19 +187,16 @@ bool CenternetDetection3DTrack::init_postprocessing(){
 
     checkCuda( cudaMallocHost(&target_coords, 4 * K *sizeof(float)) );
 
-    calibs = cv::Mat(cv::Size(4,3), CV_32F);
+    for(int bi=0; bi<nBatches; bi++) {
-    calibs.at<float>(0,0) = 633.0;
+        cv::Mat calibs_ = cv::Mat::zeros(cv::Size(4,3), CV_32F);        
-    calibs.at<float>(0,1) = 0.0;
+        if(inputCalibs.size() == 0 || inputCalibs[bi].empty()) {
-    calibs.at<float>(0,2) = 0.0; //w/2
+            calibs_.at<float>(0,0) = 633.0;
-    calibs.at<float>(0,3) = 0.0;
+            calibs_.at<float>(1,1) = 633.0;
-    calibs.at<float>(1,0) = 0.0;
+            calibs_.at<float>(2,2) = 1.0;
-    calibs.at<float>(1,1) = 633.0;
+        }
-    calibs.at<float>(1,2) = 0.0; //h/2
+        calibs_.at<float>(2,2) = 1.0;
-    calibs.at<float>(1,3) = 0.0;
+        calibs.push_back(calibs_);
-    calibs.at<float>(2,0) = 0.0;
+    }
-    calibs.at<float>(2,1) = 0.0;
-    calibs.at<float>(2,2) = 1.0;
-    calibs.at<float>(2,3) = 0.0;
 
     // Alloc array used in the kernel 
     checkCuda( cudaMalloc(&src_out, K *sizeof(float)) );
@@ -288,17 +289,25 @@ void CenternetDetection3DTrack::pre_inf(const int bi){
     checkCuda( cudaDeviceSynchronize() );
 }
 
-void CenternetDetection3DTrack::preprocess(cv::Mat &frame, const int bi){
+void CenternetDetection3DTrack::preprocess(cv::Mat &frame, const int bi, const std::vector<cv::Size>& stream_size){
     // -----------------------------------pre-process ------------------------------------------
     batchTracked.clear();
     cv::Size sz = originalSize[bi];
-    cv::Size sz_old;
     float scale = 1.0;
     float new_height = sz.height * scale;
     float new_width = sz.width * scale;
-    if(sz.height != sz_old.height && sz.width != sz_old.width){
+    if(sz.height != sz_old[bi].height && sz.width != sz_old[bi].width){
-        calibs.at<float>(0,2) = new_width / 2.0f;
+        if(inputCalibs.size() == 0 || inputCalibs[bi].empty()) {
-        calibs.at<float>(1,2) = new_height /2.0f;
+            calibs[bi].at<float>(0,2) = new_width / 2.0f;
+            calibs[bi].at<float>(1,2) = new_height /2.0f;
+        }
+        else {
+            calibs[bi].at<float>(0,0) = inputCalibs[bi].at<float>(0,0) * dim.w / stream_size[bi].width;
+            calibs[bi].at<float>(0,2) = inputCalibs[bi].at<float>(0,2) * dim.w / stream_size[bi].width;
+            calibs[bi].at<float>(1,1) = inputCalibs[bi].at<float>(1,1) * dim.h / stream_size[bi].height;
+            calibs[bi].at<float>(1,2) = inputCalibs[bi].at<float>(1,2) * dim.h / stream_size[bi].height;
+        }
+        
         float c[] = {new_width / 2.0f, new_height /2.0f};
         float s[] = {dim.w, dim.h};
         // float s = new_width >= new_height ? new_width : new_height;
@@ -324,7 +333,7 @@ void CenternetDetection3DTrack::preprocess(cv::Mat &frame, const int bi){
         trans2 = cv::getAffineTransform( dst2, src );
         trans2.convertTo(trans_out, CV_32F);
     }
-    sz_old = sz;
+    sz_old[bi] = sz;
 #ifdef OPENCV_CUDACONTRIB
     std::cout<<"OPENCV CPMTROB\n";
     cv::cuda::GpuMat im_Orig; 
@@ -358,7 +367,7 @@ void CenternetDetection3DTrack::preprocess(cv::Mat &frame, const int bi){
 #else
     std::cout<<"NO OPENCV CPMTROB\n";
     cv::Mat imageF;
-    // resize(frame, imageF, cv::Size(new_width, new_height));
+    //resize(frame, imageF, cv::Size(512, 512));
     imageF = frame;
     sz = imageF.size();
  
@@ -701,9 +710,9 @@ void CenternetDetection3DTrack::postprocess(const int bi, const bool mAP) {
         new_det_res.dim[2] = dim_[i+2*K];
         
         // unproject_2d_to_3d
-        new_det_res.z = dep[i] - calibs.at<float>(2,3);
+        new_det_res.z = dep[i] - calibs[bi].at<float>(2,3);
-        new_det_res.x = ((float)new_det_res.ct.at<float>(0,0) * dep[i] - calibs.at<float>(0,3) - calibs.at<float>(0,2) * new_det_res.z) / calibs.at<float>(0,0);
+        new_det_res.x = ((float)new_det_res.ct.at<float>(0,0) * dep[i] - calibs[bi].at<float>(0,3) - calibs[bi].at<float>(0,2) * new_det_res.z) / calibs[bi].at<float>(0,0);
-        new_det_res.y = ((float)new_det_res.ct.at<float>(0,1) * dep[i] - calibs.at<float>(1,3) - calibs.at<float>(1,2) * new_det_res.z) / calibs.at<float>(1,1) + (dim_[i] / 2);
+        new_det_res.y = ((float)new_det_res.ct.at<float>(0,1) * dep[i] - calibs[bi].at<float>(1,3) - calibs[bi].at<float>(1,2) * new_det_res.z) / calibs[bi].at<float>(1,1) + (dim_[i] / 2);
         
         // alpha2rot_y
         // idx = rot[:, 1] > rot[:, 5]
@@ -714,7 +723,7 @@ void CenternetDetection3DTrack::postprocess(const int bi, const bool mAP) {
             new_det_res.alpha = std::atan2(rot[2*K + i], rot[3*K + i]) -0.5 * M_PI;
         else
             new_det_res.alpha = std::atan2(rot[6*K + i], rot[7*K + i]) +0.5 * M_PI;
-        new_det_res.rot_y = (new_det_res.alpha + std::atan2((float)new_det_res.ct.at<float>(0,0) - calibs.at<float>(0,2), calibs.at<float>(0,0)));
+        new_det_res.rot_y = (new_det_res.alpha + std::atan2((float)new_det_res.ct.at<float>(0,0) - calibs[bi].at<float>(0,2), calibs[bi].at<float>(0,0)));
         new_det_res.ct = new_det_res.ct + new_det_res.tr;   //dest  
         det_res.push_back(new_det_res);    
 
@@ -804,7 +813,7 @@ void CenternetDetection3DTrack::draw(std::vector<cv::Mat>& frames) {
                     }
                     
                     aus.release();
-                    aus = calibs * pts3DHomo;
+                    aus = calibs[bi] * pts3DHomo;
                     std::vector<float> res_corners;
                     for(int k=0; k<8; k++) {
                         res_corners.push_back(aus.at<float>(0,k) / aus.at<float>(2,k));