[Model] add tracking trail on vis_mot

examples/vision/tracking/pptracking/cpp/infer.cc

@@ -34,27 +34,28 @@ void CpuInfer(const std::string& model_dir, const std::string& video_file) {
 
   fastdeploy::vision::MOTResult result;
   fastdeploy::vision::tracking::TrailRecorder recorder;
-  //during each prediction, data is inserted into the recorder. As the number of predictions increases,
-  //the memory will continue to grow. You can cancel the insertion through 'UnBindRecorder'.
-  int count = 0; // unbind condition
+  // during each prediction, data is inserted into the recorder. As the number of predictions increases,
+  // the memory will continue to grow. You can cancel the insertion through 'UnbindRecorder'.
+  // int count = 0; // unbind condition
   model.BindRecorder(&recorder);
   cv::Mat frame;
   cv::VideoCapture capture(video_file);
   while (capture.read(frame)) {
     if (frame.empty()) {
-        break;
+      break;
     }
     if (!model.Predict(&frame, &result)) {
-        std::cerr << "Failed to predict." << std::endl;
-        return;
+      std::cerr << "Failed to predict." << std::endl;
+      return;
     }
     // such as adding this code can cancel trail datat bind
-    if(count++ == 10) model.UnBindRecorder();
+    // if(count++ == 10) model.UnbindRecorder();
     // std::cout << result.Str() << std::endl;
     cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &recorder);
     cv::imshow("mot",out_img);
     cv::waitKey(30);
   }
+  model.UnbindRecorder();
   capture.release();
   cv::destroyAllWindows();
 }
@@ -76,27 +77,28 @@ void GpuInfer(const std::string& model_dir, const std::string& video_file) {
 
   fastdeploy::vision::MOTResult result;
   fastdeploy::vision::tracking::TrailRecorder trail_recorder;
-  //during each prediction, data is inserted into the recorder. As the number of predictions increases,
-  //the memory will continue to grow. You can cancel the insertion through 'UnBindRecorder'.
-  int count = 0; // unbind condition
+  // during each prediction, data is inserted into the recorder. As the number of predictions increases,
+  // the memory will continue to grow. You can cancel the insertion through 'UnbindRecorder'.
+  // int count = 0; // unbind condition
   model.BindRecorder(&trail_recorder);
   cv::Mat frame;
   cv::VideoCapture capture(video_file);
   while (capture.read(frame)) {
-      if (frame.empty()) {
-          break;
-      }
-      if (!model.Predict(&frame, &result)) {
-          std::cerr << "Failed to predict." << std::endl;
-          return;
-      }
-      // such as adding this code can cancel trail datat bind
-      if(count++ == 10) model.UnBindRecorder();
-      // std::cout << result.Str() << std::endl;
-      cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &trail_recorder);
-      cv::imshow("mot",out_img);
-      cv::waitKey(30);
+    if (frame.empty()) {
+      break;
+    }
+    if (!model.Predict(&frame, &result)) {
+      std::cerr << "Failed to predict." << std::endl;
+      return;
+    }
+    // such as adding this code can cancel trail datat bind
+    //if(count++ == 10) model.UnbindRecorder();
+    // std::cout << result.Str() << std::endl;
+    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &trail_recorder);
+    cv::imshow("mot",out_img);
+    cv::waitKey(30);
   }
+  model.UnbindRecorder();
   capture.release();
   cv::destroyAllWindows();
 }
@@ -120,26 +122,27 @@ void TrtInfer(const std::string& model_dir, const std::string& video_file) {
   fastdeploy::vision::MOTResult result;
   fastdeploy::vision::tracking::TrailRecorder recorder;
   //during each prediction, data is inserted into the recorder. As the number of predictions increases,
-  //the memory will continue to grow. You can cancel the insertion through 'UnBindRecorder'.
-  int count = 0; // unbind condition
+  //the memory will continue to grow. You can cancel the insertion through 'UnbindRecorder'.
+  // int count = 0; // unbind condition
   model.BindRecorder(&recorder);
   cv::Mat frame;
   cv::VideoCapture capture(video_file);
   while (capture.read(frame)) {
-      if (frame.empty()) {
-          break;
-      }
-      if (!model.Predict(&frame, &result)) {
-          std::cerr << "Failed to predict." << std::endl;
-          return;
-      }
-      // such as adding this code can cancel trail datat bind
-      if(count++ == 10) model.UnBindRecorder();
-      // std::cout << result.Str() << std::endl;
-      cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &recorder);
-      cv::imshow("mot",out_img);
-      cv::waitKey(30);
+    if (frame.empty()) {
+        break;
+    }
+    if (!model.Predict(&frame, &result)) {
+        std::cerr << "Failed to predict." << std::endl;
+        return;
+    }
+    // such as adding this code can cancel trail datat bind
+    // if(count++ == 10) model.UnbindRecorder();
+    // std::cout << result.Str() << std::endl;
+    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &recorder);
+    cv::imshow("mot",out_img);
+    cv::waitKey(30);
   }
+  model.UnbindRecorder();
   capture.release();
   cv::destroyAllWindows();
 }

examples/vision/tracking/pptracking/python/infer.py

@@ -14,7 +14,6 @@
 
 import fastdeploy as fd
 import cv2
-import time
 import os
 
 
@@ -61,10 +60,10 @@ def build_option(args):
     model_file, params_file, config_file, runtime_option=runtime_option)
 
 # 初始化轨迹记录器
-trail_recorder = fd.vision.tracking.TrailRecorder()
+recorder = fd.vision.tracking.TrailRecorder()
 # 绑定记录器 注意：每次预测时，往trail_recorder里面插入数据，随着预测次数的增加，内存会不断地增长，
-# 可以通过unbind_trail_recorders()方法来解除绑定
-model.bind_trail_recorders(trail_recorder)
+# 可以通过unbind_recorder()方法来解除绑定
+model.bind_recorder(recorder)
 # 预测图片分割结果
 cap = cv2.VideoCapture(args.video)
 # count = 0
@@ -75,10 +74,11 @@ def build_option(args):
     result = model.predict(frame)
     # count += 1
     # if count == 10:
-    #     model.unbind_trail_recorders()
-    img = fd.vision.vis_mot(frame, result, 0.0)
+    #     model.unbind_recorder()
+    img = fd.vision.vis_mot(frame, result, 0.0, recorder)
     cv2.imshow("video", img)
     if cv2.waitKey(30) == ord("q"):
         break
+model.unbind_recorder()
 cap.release()
 cv2.destroyAllWindows()

fastdeploy/vision/tracking/pptracking/model.cc

@@ -261,7 +261,6 @@ bool PPTracking::Postprocess(std::vector<FDTensor>& infer_result, MOTResult *res
     result->boxes.push_back(box);
     result->ids.push_back(1);
     result->scores.push_back(*dets.ptr<float>(0, 4));
-
   } else {
     std::vector<Track>::iterator titer;
     for (titer = tracks.begin(); titer != tracks.end(); ++titer) {
@@ -288,14 +287,7 @@ bool PPTracking::Postprocess(std::vector<FDTensor>& infer_result, MOTResult *res
     float center_x = (result->boxes[i][0] + result->boxes[i][2]) / 2;
     float center_y = (result->boxes[i][1] + result->boxes[i][3]) / 2;
     int id = result->ids[i];
-    auto iter = recorder_->records.find(id);
-    if (iter != recorder_->records.end()) {
-      auto trail = recorder_->records[id];
-      trail.push_back({int(center_x), int(center_y)});
-      recorder_->Add(id, trail);
-    } else {
-      recorder_->Add(id, {{int(center_x), int(center_y)}});
-    }
+    recorder_->Add(id,{int(center_x), int(center_y)});
   }
   return true;
 }
@@ -306,11 +298,11 @@ void PPTracking::BindRecorder(TrailRecorder* recorder){
     is_record_trail_ = true;
 }
 
-void PPTracking::UnBindRecorder(){
+void PPTracking::UnbindRecorder(){
 
     is_record_trail_ = false;
     std::map<int, std::vector<std::array<int, 2>>>::iterator iter;
-    for(iter=recorder_->records.begin();iter!=recorder_->records.end();iter++){
+    for(iter = recorder_->records.begin(); iter != recorder_->records.end(); iter++){
       iter->second.clear();
       iter->second.shrink_to_fit();
     }

fastdeploy/vision/tracking/pptracking/model.h

@@ -25,12 +25,18 @@ namespace vision {
 namespace tracking {
 struct TrailRecorder{
   std::map<int, std::vector<std::array<int, 2>>> records;
-  void Add(int id, const std::vector<std::array<int, 2>>& record);
+  void Add(int id, const std::array<int, 2>& record);
 };
 
-inline void TrailRecorder::Add(int id, const
-                               std::vector<std::array<int, 2>> &record) {
-  records[id] = record;
+inline void TrailRecorder::Add(int id, const std::array<int, 2>& record) {
+  auto iter = records.find(id);
+  if (iter != records.end()) {
+    auto trail = records[id];
+    trail.push_back(record);
+    records[id] = trail;
+  } else {
+    records[id] = {record};
+  }
 }
 
 class FASTDEPLOY_DECL PPTracking: public FastDeployModel {
@@ -66,7 +72,7 @@ class FASTDEPLOY_DECL PPTracking: public FastDeployModel {
   void BindRecorder(TrailRecorder* recorder);
   /** \brief cancel binding and clear trail information
    */
-  void UnBindRecorder();
+  void UnbindRecorder();
 
 
  private:

fastdeploy/vision/tracking/pptracking/pptracking_pybind.cc

@@ -32,7 +32,7 @@ void BindPPTracking(pybind11::module &m) {
              self.Predict(&mat, res);
              return res;
          })
-    .def("bind_trail_recorders", &vision::tracking::PPTracking::BindRecorder)
-    .def("unbind_trail_recorders", &vision::tracking::PPTracking::UnBindRecorder);
+    .def("bind_recorder", &vision::tracking::PPTracking::BindRecorder)
+    .def("unbind_recorder", &vision::tracking::PPTracking::UnbindRecorder);
 }
 }  // namespace fastdeploy

python/fastdeploy/vision/tracking/pptracking/__init__.py

@@ -15,11 +15,13 @@
 from __future__ import absolute_import
 from .... import FastDeployModel, ModelFormat
 from .... import c_lib_wrap as C
+import logging
 
-
-class TrailRecorder(C.vision.tracking.TrailRecorder):
-    def __init__(self):
-        super(TrailRecorder, self).__init__()
+try:
+    TrailRecorder = C.vision.tracking.TrailRecorder
+except Exception as e:
+    logging.warning("something was wrong, detail:" + str(e) +
+                    "so 'TrailRecorder' can not be loaded")
 
 
 class PPTracking(FastDeployModel):
@@ -54,8 +56,17 @@ def predict(self, input_image):
         assert input_image is not None, "The input image data is None."
         return self._model.predict(input_image)
 
-    def bind_trail_recorders(self, val):
-        self._model.bind_trail_recorders(val)
+    def bind_recorder(self, val):
+        """ Binding tracking trail
+
+        :param val: (TrailRecorder) trail recorder, which is contained object's id and center point sequence
+        :return: None
+        """
+        self._model.bind_recorder(val)
 
-    def unbind_trail_recorders(self):
-        self._model.unbind_trail_recorders()
+    def unbind_recorder(self):
+        """ cancel binding of tracking trail
+
+        :return:
+        """
+        self._model.unbind_recorder()