提交gpu代码

BaseTracker/BaseTrack.cpp

@@ -81,3 +81,29 @@ int BaseTrack::GetMeasureData(std::vector<float>& data)
         data.push_back(kf_->x_[i]);
     return 0;
 }
+
+int BaseTrack::GetStatesNum()
+{
+    if (kf_ == nullptr)
+        return 0;
+    return kf_->num_states_;
+}
+int BaseTrack::GetObsNum()
+{
+    if (kf_ == nullptr)
+        return 0;
+    return kf_->num_obs_;
+}
+
+float* BaseTrack::GetStatesXPtr()
+{
+    if (kf_ == nullptr)
+        return nullptr;
+    return kf_->x_.data();
+}
+float* BaseTrack::GetPredictPtr()
+{
+    if (kf_ == nullptr)
+        return nullptr;
+    return kf_->P_.data();
+}
\ No newline at end of file

BaseTracker/BaseTrack.h

@@ -27,6 +27,12 @@ public:
 
     virtual double CalculateIou(const std::vector<float>& data) = 0;
 
+    int GetStatesNum();
+    int GetObsNum();
+
+    float* GetStatesXPtr();
+    float* GetPredictPtr();
+
     int coast_cycles_ = 0, hit_streak_ = 0;
 
     int m_num_states = 0;

BaseTracker/BaseTracker.h

@@ -77,8 +77,35 @@ int BaseTracker<T>::Run(const std::vector<std::vector<float> >& detections, std:
     }
     else
     {
-        //float* Z[matched.size()] = {};
-        std::shared_ptr<float*> Z = std::shared_ptr<float*>(new float*[10], [](float** p) {if (p) delete[] p; p = nullptr; });
+        int bs = matched.size();
+        int ns = 0;
+        int no = detections.size() > 0 ? detections[0].size():0;
+        std::shared_ptr<float> Z = std::shared_ptr<float>(new float[bs * no], [](float* p) {if (p) delete[] p; p = nullptr; });
+        std::shared_ptr<float*> X = std::shared_ptr<float*>(new float* [bs], [](float** p) {if (p) delete[] p; p = nullptr; });
+        std::shared_ptr<float*> P = std::shared_ptr<float*>(new float* [bs], [](float** p) {if (p) delete[] p; p = nullptr; });
+        std::shared_ptr<float> HX = std::shared_ptr<float>(new float[bs * no], [](float* p) {if (p) delete[] p; p = nullptr; });
+        int bs_i = 0;
+        for (const auto& match : matched)
+        {
+            const auto& id = match.first;
+            float* ptr_Z = Z.get() + bs_i * no;
+            for (int i = 0; i < no; i++)
+            {
+                ptr_Z[i] = detections[match.second][i];
+            }
+            X.get()[bs_i] = m_tracker[id]->GetStatesXPtr();
+            P.get()[bs_i] = m_tracker[id]->GetPredictPtr();
+            float* ptr_HX = HX.get() + bs_i * no;
+            memcpy(ptr_HX, m_tracker[id]->GetStatesXPtr(), no);
+            if (ns == 0)
+            {
+                ns = m_tracker[id]->GetStatesNum();
+            }
+            bs_i++;
+            detectionsId[match.second] = id;
+            updateId[id] = match.second;
+        }
+        //kalman_update_batch(Z.get(), X.get(), P.get(), HX.get(), bs, ns, no);
     }
 
     /*** Create new tracks for unmatched detections ***/

BaseTracker/kf_gpu/kalman_update_batch_online.cu

@@ -309,10 +309,10 @@ void map_kalman_update_batch( 	pybind11::array_t<float> Z,
 
   MAKE SURE ALL INPUTS ARE TWO-DIM NUMPY ARRAY
   */
-void kalman_update_batch(float** Z,// measurement size = bs * no
+void kalman_update_batch(float* Z,// measurement size = bs * no
 	float** X,     // in-place update states  size = bs * ns
 	float** P,     // in-place update predict size = bs * ns * ns
-	float** HX,	   // H*X  size = bs * no
+	float* HX,	   // H*X  size = bs * no
 	const int bs,
 	const int ns,  //ns = 10
 	const int no   // no = 7
@@ -344,13 +344,14 @@ void kalman_update_batch(float** Z,// measurement size = bs * no
 	GPU_CHECK(cudaMalloc(&device_X, size_XX));
 	GPU_CHECK(cudaMalloc(&device_P, size_PP));
 	GPU_CHECK(cudaMalloc(&device_HX, size_HXX));
+	GPU_CHECK(cudaMemcpy(device_Z , Z, size_ZZ, cudaMemcpyHostToDevice));
 	for (int i = 0; i < bs; i++)
 	{
-		GPU_CHECK(cudaMemcpy(device_Z + i*no, Z[i], no * sizeof(float), cudaMemcpyHostToDevice));
 		GPU_CHECK(cudaMemcpy(device_X + i*ns, X[i], ns * sizeof(float), cudaMemcpyHostToDevice));
 		GPU_CHECK(cudaMemcpy(device_P + i*ns*ns, P[i], ns*ns * sizeof(float), cudaMemcpyHostToDevice));
-		GPU_CHECK(cudaMemcpy(device_HX + i*no, HX[i], no * sizeof(float), cudaMemcpyHostToDevice));
 	}
+	GPU_CHECK(cudaMemcpy(device_HX, HX, size_HXX, cudaMemcpyHostToDevice));
+
 	kalmanUpdateLauncher_batch(device_Z, device_X, device_P, device_HX, bs, ns, no);
 
 	for (int i = 0; i < bs; i++)

BaseTracker/kf_gpu/kalman_update_batch_online.h

@@ -3,10 +3,10 @@
 
 
 
-void kalman_update_batch(float** Z,// measurement size = bs * no
+void kalman_update_batch(float* Z,// measurement size = bs * no
 	float** X,     // in-place update states  size = bs * ns
 	float** P,     // in-place update predict size = bs * ns * ns
-	float** HX,	   // H*X  size = bs * no
+	float* HX,	   // H*X  size = bs * no
 	const int bs,
 	const int ns,  //ns = 10
 	const int no   // no = 7