From ff6e0e010adc120bd010c605f82355111758dd08 Mon Sep 17 00:00:00 2001
From: Davide Sapienza <sapienza.dav@gmail.com>
Date: Fri, 30 Apr 2021 22:26:33 +0200
Subject: [PATCH] Fix a bug with batch > 1

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
---
 include/tkDNN/CenternetDetection3DTrack.h |   9 +-
 src/CenternetDetection3DTrack.cpp         | 184 +++++++++++-----------
 2 files changed, 99 insertions(+), 94 deletions(-)

diff --git a/include/tkDNN/CenternetDetection3DTrack.h b/include/tkDNN/CenternetDetection3DTrack.h
index 451bf6e..258fe9d 100644
--- a/include/tkDNN/CenternetDetection3DTrack.h
+++ b/include/tkDNN/CenternetDetection3DTrack.h
@@ -148,10 +148,9 @@ private:
     std::vector<struct detectionRes> det_res;
     int count_det;
     //tracks
-    std::vector<struct trackingRes> tr_res;
-    std::vector<std::vector<struct trackingRes>> batchTracked;
-    int count_tr;
-    int track_id=0;
+    std::vector<std::vector<struct trackingRes>> tr_res;
+    std::vector<int> count_tr;
+    std::vector<int> track_id;
   
     
     bool init_preprocessing();
@@ -161,7 +160,7 @@ private:
     void pre_inf(const int bi);
     void _get_additional_inputs();
     cv::Mat transform_preds_with_trans(float x1, float x2);
-    void tracking();
+    void tracking(const int bi);
 
 public:
     tk::dnn::Network *pre_phase_net = nullptr;
diff --git a/src/CenternetDetection3DTrack.cpp b/src/CenternetDetection3DTrack.cpp
index 02db674..bdc1c59 100644
--- a/src/CenternetDetection3DTrack.cpp
+++ b/src/CenternetDetection3DTrack.cpp
@@ -17,8 +17,6 @@ bool CenternetDetection3DTrack::init(const std::string& tensor_path, const int n
     init_pre_inf();
     init_postprocessing();
     init_visualization(n_classes);
-    
-    count_tr = 0;
 }
 
 bool CenternetDetection3DTrack::init_preprocessing(){
@@ -201,6 +199,11 @@ bool CenternetDetection3DTrack::init_postprocessing(){
     // Alloc array used in the kernel 
     checkCuda( cudaMalloc(&src_out, K *sizeof(float)) );
     checkCuda( cudaMalloc(&ids_out, K *sizeof(int)) );
+
+    for(int bi=0; bi<nBatches; bi++){
+        track_id.push_back(0);
+        count_tr.push_back(0);
+    }
 }
 
 bool CenternetDetection3DTrack::init_visualization(const int n_classes){
@@ -291,7 +294,6 @@ void CenternetDetection3DTrack::pre_inf(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];
     float scale = 1.0;
     float new_height = sz.height * scale;
@@ -414,8 +416,7 @@ cv::Mat CenternetDetection3DTrack::transform_preds_with_trans(float x1, float x2
     return trans_out * target_coords;
 }
 
-void CenternetDetection3DTrack::tracking(){
-   
+void CenternetDetection3DTrack::tracking(const int bi) {
     float item_size[count_det];
     int item_cl[count_det];
     float dets[2*count_det];
@@ -427,44 +428,44 @@ void CenternetDetection3DTrack::tracking(){
         dets[i*2+1] = det_res[i].ct.at<float>(0,1);
     }
     
-    float track_size[count_tr];
-    int track_cl[count_tr];
-    float tracks[2*count_tr];
-    for(int i=0; i<count_tr; i++){
-        track_size[i] = (tr_res[i].det_res.bb1.at<float>(0,0) - tr_res[i].det_res.bb0.at<float>(0,0)) * 
-                        (tr_res[i].det_res.bb1.at<float>(0,1) - tr_res[i].det_res.bb0.at<float>(0,1));
-        track_cl[i] = tr_res[i].det_res.cl;
-        tracks[i*2] = tr_res[i].det_res.ct.at<float>(0,0);
-        tracks[i*2+1] = tr_res[i].det_res.ct.at<float>(0,1);
+    float track_size[count_tr[bi]];
+    int track_cl[count_tr[bi]];
+    float tracks[2*count_tr[bi]];
+    for(int i=0; i<count_tr[bi]; i++){
+        track_size[i] = (tr_res[bi][i].det_res.bb1.at<float>(0,0) - tr_res[bi][i].det_res.bb0.at<float>(0,0)) * 
+                        (tr_res[bi][i].det_res.bb1.at<float>(0,1) - tr_res[bi][i].det_res.bb0.at<float>(0,1));
+        track_cl[i] = tr_res[bi][i].det_res.cl;
+        tracks[i*2] = tr_res[bi][i].det_res.ct.at<float>(0,0);
+        tracks[i*2+1] = tr_res[bi][i].det_res.ct.at<float>(0,1);
     } 
-    float dist[count_tr*count_det];
+    float dist[count_tr[bi]*count_det];
     bool invalid;
-    for(int i=0; i<count_tr; i++){
+    for(int i=0; i<count_tr[bi]; i++){
         for(int j=0; j<count_det; j++){
-            dist[j*count_tr+i] =  pow((tracks[i*2] - dets[j*2]), 2) + 
+            dist[j*count_tr[bi]+i] =  pow((tracks[i*2] - dets[j*2]), 2) + 
                             pow((tracks[i*2+1] - dets[j*2+1]), 2);
-            invalid = dist[j*count_tr+i] > track_size[i] || dist[j*count_tr+i] > item_size[j] || item_cl[j] != track_cl[i];
-            dist[j*count_tr+i] = dist[j*count_tr+i] + invalid * (1 << 18);
+            invalid = dist[j*count_tr[bi]+i] > track_size[i] || dist[j*count_tr[bi]+i] > item_size[j] || item_cl[j] != track_cl[i];
+            dist[j*count_tr[bi]+i] = dist[j*count_tr[bi]+i] + invalid * (1 << 18);
         }
     }
-    int matched_indices[2*count_tr];
+    int matched_indices[2*count_tr[bi]];
     float min_tr;
     int min_idtr=-1;
-    for(int i=0; i<count_tr; i++) {
+    for(int i=0; i<count_tr[bi]; i++) {
         matched_indices[i*2] = -1;
         matched_indices[i*2+1] = -1;
     }
     for(int i=0; i<count_det; i++){
         min_tr=(1 << 18);
-        for(int j=0; j<count_tr; j++){
-            if(dist[i*count_tr+j]<min_tr) {
-                min_tr = dist[i*count_tr+j];
+        for(int j=0; j<count_tr[bi]; j++){
+            if(dist[i*count_tr[bi]+j]<min_tr) {
+                min_tr = dist[i*count_tr[bi]+j];
                 min_idtr = j;
             }
         }
         if(min_tr < (1<<16)) {
             for(int j=0; j<count_det; j++){
-                dist[j*count_tr+min_idtr] = (1 << 18);
+                dist[j*count_tr[bi]+min_idtr] = (1 << 18);
             }
             matched_indices[2*min_idtr] = min_idtr;
             matched_indices[2*min_idtr+1] = i;
@@ -474,10 +475,10 @@ void CenternetDetection3DTrack::tracking(){
     bool unmatched_dets[count_det];
     for(int i=0; i<count_det; i++)
         unmatched_dets[i] = false;
-    bool unmatched_tracks[count_tr];
-    for(int i=0; i<count_tr; i++) 
+    bool unmatched_tracks[count_tr[bi]];
+    for(int i=0; i<count_tr[bi]; i++) 
         unmatched_tracks[i] = false;
-    for(int i=0; i<count_tr; i++) {
+    for(int i=0; i<count_tr[bi]; i++) {
         if(matched_indices[2*i] != -1)
             unmatched_tracks[matched_indices[2*i]]=true;
         
@@ -486,84 +487,83 @@ void CenternetDetection3DTrack::tracking(){
     }
 
     //match
-    for(int i=0; i<count_tr; i++) {
+    for(int i=0; i<count_tr[bi]; i++) {
         if(matched_indices[2*i+1] != -1 && matched_indices[2*i] != -1) { //second condition is optional
             int tr_id = matched_indices[2*i];
             int d_id = matched_indices[2*i+1];
             
             // tr_res[tr_id].det_res = det_res[d_id];
-            tr_res[tr_id].det_res.score = det_res[d_id].score;
-            tr_res[tr_id].det_res.cl = det_res[d_id].cl;
-            tr_res[tr_id].det_res.ct = det_res[d_id].ct;
-            tr_res[tr_id].det_res.tr = det_res[d_id].tr;
-            tr_res[tr_id].det_res.bb0 = det_res[d_id].bb0;
-            tr_res[tr_id].det_res.bb1 = det_res[d_id].bb1;
-            tr_res[tr_id].det_res.dep = det_res[d_id].dep;
-            tr_res[tr_id].det_res.dim[0] = det_res[d_id].dim[0];
-            tr_res[tr_id].det_res.dim[1] = det_res[d_id].dim[1];
-            tr_res[tr_id].det_res.dim[2] = det_res[d_id].dim[2];
-            tr_res[tr_id].det_res.alpha = det_res[d_id].alpha;
-            tr_res[tr_id].det_res.x = det_res[d_id].x;
-            tr_res[tr_id].det_res.y = det_res[d_id].y;
-            tr_res[tr_id].det_res.z = det_res[d_id].z;
-            tr_res[tr_id].det_res.rot_y = det_res[d_id].rot_y;
-            // tr_res[matched_indices[2*i]].tracking_id = ; is the same 
-            // tr_res[matched_indices[2*i]].color = ; is the same
-            tr_res[tr_id].age = 1; 
-            tr_res[tr_id].active = tr_res[tr_id].active+1;
+            tr_res[bi][tr_id].det_res.score = det_res[d_id].score;
+            tr_res[bi][tr_id].det_res.cl = det_res[d_id].cl;
+            tr_res[bi][tr_id].det_res.ct = det_res[d_id].ct;
+            tr_res[bi][tr_id].det_res.tr = det_res[d_id].tr;
+            tr_res[bi][tr_id].det_res.bb0 = det_res[d_id].bb0;
+            tr_res[bi][tr_id].det_res.bb1 = det_res[d_id].bb1;
+            tr_res[bi][tr_id].det_res.dep = det_res[d_id].dep;
+            tr_res[bi][tr_id].det_res.dim[0] = det_res[d_id].dim[0];
+            tr_res[bi][tr_id].det_res.dim[1] = det_res[d_id].dim[1];
+            tr_res[bi][tr_id].det_res.dim[2] = det_res[d_id].dim[2];
+            tr_res[bi][tr_id].det_res.alpha = det_res[d_id].alpha;
+            tr_res[bi][tr_id].det_res.x = det_res[d_id].x;
+            tr_res[bi][tr_id].det_res.y = det_res[d_id].y;
+            tr_res[bi][tr_id].det_res.z = det_res[d_id].z;
+            tr_res[bi][tr_id].det_res.rot_y = det_res[d_id].rot_y;
+            // tr_res[bi][matched_indices[2*i]].tracking_id = ; is the same 
+            // tr_res[bi][matched_indices[2*i]].color = ; is the same
+            tr_res[bi][tr_id].age = 1; 
+            tr_res[bi][tr_id].active = tr_res[bi][tr_id].active+1;
         }
     }
     //delete target umatched track
     int new_count_tr = 0;
-    for(int i=0; i<count_tr; i++) {
+    for(int i=0; i<count_tr[bi]; i++) {
         if(unmatched_tracks[i]) 
             new_count_tr++;
     }
-    if(new_count_tr == 0 && count_tr != 0) {        //reset
-        tr_res.clear();
-        count_tr = 0;
+    if(new_count_tr == 0 && count_tr[bi] != 0) {        //reset
+        tr_res[bi].clear();
+        count_tr[bi] = 0;
     }
-
-    int old_count_tr = count_tr;
-    if(count_tr != 0 && new_count_tr != count_tr) {
+    int old_count_tr = count_tr[bi];
+    if(count_tr[bi] != 0 && new_count_tr != count_tr[bi]) {
         std::vector<struct trackingRes> new_tr_res;
         int id_new_tr=0;
-        for(int i=0; i<count_tr; i++) {
+        for(int i=0; i<count_tr[bi]; i++) {
             if(unmatched_tracks[i]) {
                 struct trackingRes new_tr_res_;
                 // new_tr_res_new_det_res.det_res = tr_res[i].det_res;
-                new_tr_res_.det_res.score = tr_res[i].det_res.score;
-                new_tr_res_.det_res.cl = tr_res[i].det_res.cl;
-                new_tr_res_.det_res.ct = tr_res[i].det_res.ct;
-                new_tr_res_.det_res.tr = tr_res[i].det_res.tr;
-                new_tr_res_.det_res.bb0 = tr_res[i].det_res.bb0;
-                new_tr_res_.det_res.bb1 = tr_res[i].det_res.bb1;
-                new_tr_res_.det_res.dep = tr_res[i].det_res.dep;
-                new_tr_res_.det_res.dim[0] = tr_res[i].det_res.dim[0];
-                new_tr_res_.det_res.dim[1] = tr_res[i].det_res.dim[1];
-                new_tr_res_.det_res.dim[2] = tr_res[i].det_res.dim[2];
-                new_tr_res_.det_res.alpha = tr_res[i].det_res.alpha;
-                new_tr_res_.det_res.x = tr_res[i].det_res.x;
-                new_tr_res_.det_res.y = tr_res[i].det_res.y;
-                new_tr_res_.det_res.z = tr_res[i].det_res.z;
-                new_tr_res_.det_res.rot_y = tr_res[i].det_res.rot_y;
-                new_tr_res_.tracking_id = tr_res[i].tracking_id;
-                new_tr_res_.age = tr_res[i].age;
-                new_tr_res_.active = tr_res[i].active;
-                new_tr_res_.color = tr_res[i].color;
+                new_tr_res_.det_res.score = tr_res[bi][i].det_res.score;
+                new_tr_res_.det_res.cl = tr_res[bi][i].det_res.cl;
+                new_tr_res_.det_res.ct = tr_res[bi][i].det_res.ct;
+                new_tr_res_.det_res.tr = tr_res[bi][i].det_res.tr;
+                new_tr_res_.det_res.bb0 = tr_res[bi][i].det_res.bb0;
+                new_tr_res_.det_res.bb1 = tr_res[bi][i].det_res.bb1;
+                new_tr_res_.det_res.dep = tr_res[bi][i].det_res.dep;
+                new_tr_res_.det_res.dim[0] = tr_res[bi][i].det_res.dim[0];
+                new_tr_res_.det_res.dim[1] = tr_res[bi][i].det_res.dim[1];
+                new_tr_res_.det_res.dim[2] = tr_res[bi][i].det_res.dim[2];
+                new_tr_res_.det_res.alpha = tr_res[bi][i].det_res.alpha;
+                new_tr_res_.det_res.x = tr_res[bi][i].det_res.x;
+                new_tr_res_.det_res.y = tr_res[bi][i].det_res.y;
+                new_tr_res_.det_res.z = tr_res[bi][i].det_res.z;
+                new_tr_res_.det_res.rot_y = tr_res[bi][i].det_res.rot_y;
+                new_tr_res_.tracking_id = tr_res[bi][i].tracking_id;
+                new_tr_res_.age = tr_res[bi][i].age;
+                new_tr_res_.active = tr_res[bi][i].active;
+                new_tr_res_.color = tr_res[bi][i].color;
                 id_new_tr++;
                 new_tr_res.push_back(new_tr_res_);
             }
         }
         
-        if(count_tr) {
-            tr_res.clear();
+        if(count_tr[bi]) {
+            tr_res[bi].clear();
         }
-        count_tr = new_count_tr;
-        tr_res=new_tr_res;
+        count_tr[bi] = new_count_tr;
+        tr_res[bi]=new_tr_res;
     }
     
-    int count_tr_ = count_tr;
+    int count_tr_ = count_tr[bi];
     for(int i=0; i<count_det; i++) {
         if((!unmatched_dets[i]) && det_res[i].score > new_thresh) {
             count_tr_ ++;
@@ -583,17 +583,24 @@ void CenternetDetection3DTrack::tracking(){
             new_tr_res_.det_res.y = det_res[i].y;
             new_tr_res_.det_res.z = det_res[i].z;
             new_tr_res_.det_res.rot_y = det_res[i].rot_y;
-            new_tr_res_.tracking_id = track_id++; 
+            new_tr_res_.tracking_id = track_id[bi]++; 
             new_tr_res_.age = 1; 
             new_tr_res_.active = 1;
             new_tr_res_.color = rand() % 256;
-            tr_res.push_back(new_tr_res_);
+            if(tr_res.size() <= bi) {
+                std::vector<struct trackingRes> v_new_tr_res_;
+                v_new_tr_res_.push_back(new_tr_res_);    
+                tr_res.push_back(v_new_tr_res_);
+            }
+            else
+                tr_res[bi].push_back(new_tr_res_);
         }
     }
-    count_tr = count_tr_;
+
+    count_tr[bi] = count_tr_;
     
-    if(track_id==1000)
-        track_id=0;
+    if(track_id[bi]==1000)
+        track_id[bi]=0;
     det_res.clear();
     
 }
@@ -729,8 +736,7 @@ void CenternetDetection3DTrack::postprocess(const int bi, const bool mAP) {
 
     }    
     // track step
-    tracking();
-    batchTracked.push_back(tr_res);
+    tracking(bi);
 }
 
 void CenternetDetection3DTrack::draw(std::vector<cv::Mat>& frames) {
@@ -743,8 +749,8 @@ void CenternetDetection3DTrack::draw(std::vector<cv::Mat>& frames) {
     int thickness = 2;
     for(int bi=0; bi<frames.size(); ++bi) {
         // draw dets
-        for(int i=0; i<batchTracked[bi].size(); i++) {
-            t = batchTracked[bi][i];
+        for(int i=0; i<tr_res[bi].size(); i++) {
+            t = tr_res[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);