[Model] add tracking trail on vis_mot (PaddlePaddle#461)

* add override mark * delete some * recovery * recovery * add tracking * add tracking py_bind and example * add pptracking * add pptracking * iomanip head file * add opencv_video lib * add python libs package Signed-off-by: ChaoII <[email protected]> * complete comments Signed-off-by: ChaoII <[email protected]> * add jdeTracker_ member variable Signed-off-by: ChaoII <[email protected]> * add 'FASTDEPLOY_DECL' macro Signed-off-by: ChaoII <[email protected]> * remove kwargs params Signed-off-by: ChaoII <[email protected]> * [Doc]update pptracking docs * delete 'ENABLE_PADDLE_FRONTEND' switch * add pptracking unit test * update pptracking unit test Signed-off-by: ChaoII <[email protected]> * modify test video file path and remove trt test * update unit test model url * remove 'FASTDEPLOY_DECL' macro Signed-off-by: ChaoII <[email protected]> * fix build python packages about pptracking on win32 Signed-off-by: ChaoII <[email protected]> * update comment Signed-off-by: ChaoII <[email protected]> * add pptracking model explain Signed-off-by: ChaoII <[email protected]> * add tracking trail on vis_mot * add tracking trail * modify code for some suggestion * remove unused import * fix import bug Signed-off-by: ChaoII <[email protected]> Co-authored-by: Jason <[email protected]>
fasby45 · Nov 3, 2022 · 22d60fd · 22d60fd
1 parent 328212f
commit 22d60fd
Show file tree

Hide file tree

Showing 16 changed files with 208 additions and 116 deletions.
diff --git a/docs/api/vision_results/mot_result.md b/docs/api/vision_results/mot_result.md
@@ -37,4 +37,3 @@ fastdeploy.vision.MOTResult
 - **ids**(list of list(float)):成员变量，表示单帧画面中所有目标的id，其元素个数与`boxes`一致
 - **scores**(list of float): 成员变量，表示单帧画面检测出来的所有目标置信度
 - **class_ids**(list of int): 成员变量，表示单帧画面出来的所有目标类别
-
diff --git a/examples/vision/README.md b/examples/vision/README.md
@@ -2,16 +2,18 @@
 
 本目录下提供了各类视觉模型的部署，主要涵盖以下任务类型
 
-| 任务类型           | 说明                                  | 预测结果结构体                                                                          |
-|:-------------- |:----------------------------------- |:-------------------------------------------------------------------------------- |
-| Detection      | 目标检测，输入图像，检测图像中物体位置，并返回检测框坐标及类别和置信度 | [DetectionResult](../../docs/api/vision_results/detection_result.md)       |
-| Segmentation   | 语义分割，输入图像，给出图像中每个像素的分类及置信度          | [SegmentationResult](../../docs/api/vision_results/segmentation_result.md) |
-| Classification | 图像分类，输入图像，给出图像的分类结果和置信度             | [ClassifyResult](../../docs/api/vision_results/classification_result.md)   |
-| FaceDetection | 人脸检测，输入图像，检测图像中人脸位置，并返回检测框坐标及人脸关键点             | [FaceDetectionResult](../../docs/api/vision_results/face_detection_result.md)   |
-| KeypointDetection   | 关键点检测，输入图像，返回图像中人物行为的各个关键点坐标和置信度         | [KeyPointDetectionResult](../../docs/api/vision_results/keypointdetection_result.md) |
-| FaceRecognition | 人脸识别，输入图像，返回可用于相似度计算的人脸特征的embedding            | [FaceRecognitionResult](../../docs/api/vision_results/face_recognition_result.md)   |
-| Matting | 抠图，输入图像，返回图片的前景每个像素点的Alpha值            | [MattingResult](../../docs/api/vision_results/matting_result.md)   |
-| OCR | 文本框检测，分类，文本框内容识别，输入图像，返回文本框坐标，文本框的方向类别以及框内的文本内容            | [OCRResult](../../docs/api/vision_results/ocr_result.md)   |
+| 任务类型              | 说明                                              | 预测结果结构体                                                                              |
+|:------------------|:------------------------------------------------|:-------------------------------------------------------------------------------------|
+| Detection         | 目标检测，输入图像，检测图像中物体位置，并返回检测框坐标及类别和置信度             | [DetectionResult](../../docs/api/vision_results/detection_result.md)                 |
+| Segmentation      | 语义分割，输入图像，给出图像中每个像素的分类及置信度                      | [SegmentationResult](../../docs/api/vision_results/segmentation_result.md)           |
+| Classification    | 图像分类，输入图像，给出图像的分类结果和置信度                         | [ClassifyResult](../../docs/api/vision_results/classification_result.md)             |
+| FaceDetection     | 人脸检测，输入图像，检测图像中人脸位置，并返回检测框坐标及人脸关键点              | [FaceDetectionResult](../../docs/api/vision_results/face_detection_result.md)        |
+| KeypointDetection | 关键点检测，输入图像，返回图像中人物行为的各个关键点坐标和置信度                | [KeyPointDetectionResult](../../docs/api/vision_results/keypointdetection_result.md) |
+| FaceRecognition   | 人脸识别，输入图像，返回可用于相似度计算的人脸特征的embedding             | [FaceRecognitionResult](../../docs/api/vision_results/face_recognition_result.md)    |
+| Matting           | 抠图，输入图像，返回图片的前景每个像素点的Alpha值                     | [MattingResult](../../docs/api/vision_results/matting_result.md)                     |
+| OCR               | 文本框检测，分类，文本框内容识别，输入图像，返回文本框坐标，文本框的方向类别以及框内的文本内容 | [OCRResult](../../docs/api/vision_results/ocr_result.md)                             |
+| MOT               | 多目标跟踪，输入图像，检测图像中物体位置，并返回检测框坐标，对象id及类别置信度        | [MOTResult](../../docs/api/vision_results/mot_result.md)                             |
+
 ## FastDeploy API设计
 
 视觉模型具有较有统一任务范式，在设计API时（包括C++/Python），FastDeploy将视觉模型的部署拆分为四个步骤

diff --git a/examples/vision/tracking/pptracking/cpp/infer.cc b/examples/vision/tracking/pptracking/cpp/infer.cc
@@ -33,25 +33,29 @@ 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
+  model.BindRecorder(&recorder);
   cv::Mat frame;
-  int frame_id=0;
   cv::VideoCapture capture(video_file);
-  // according to the time of prediction to calculate fps
-  float fps= 0.0f;
   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();
     // std::cout << result.Str() << std::endl;
-    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, fps , frame_id);
+    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &recorder);
     cv::imshow("mot",out_img);
     cv::waitKey(30);
-    frame_id++;
   }
+  model.UnbindRecorder();
   capture.release();
   cv::destroyAllWindows();
 }
@@ -72,25 +76,29 @@ 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
+  model.BindRecorder(&trail_recorder);
   cv::Mat frame;
-  int frame_id=0;
   cv::VideoCapture capture(video_file);
-  // according to the time of prediction to calculate fps
-  float fps= 0.0f;
   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();
     // std::cout << result.Str() << std::endl;
-    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, fps , frame_id);
+    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &trail_recorder);
     cv::imshow("mot",out_img);
     cv::waitKey(30);
-    frame_id++;
   }
+  model.UnbindRecorder();
   capture.release();
   cv::destroyAllWindows();
 }
@@ -112,11 +120,13 @@ 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
+  model.BindRecorder(&recorder);
   cv::Mat frame;
-  int frame_id=0;
   cv::VideoCapture capture(video_file);
-  // according to the time of prediction to calculate fps
-  float fps= 0.0f;
   while (capture.read(frame)) {
     if (frame.empty()) {
         break;
@@ -125,12 +135,14 @@ void TrtInfer(const std::string& model_dir, const std::string& video_file) {
         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, fps , frame_id);
+    cv::Mat out_img = fastdeploy::vision::VisMOT(frame, result, 0.0, &recorder);
     cv::imshow("mot",out_img);
     cv::waitKey(30);
-    frame_id++;
   }
+  model.UnbindRecorder();
   capture.release();
   cv::destroyAllWindows();
 }

diff --git a/examples/vision/tracking/pptracking/python/infer.py b/examples/vision/tracking/pptracking/python/infer.py
@@ -14,7 +14,6 @@
 
 import fastdeploy as fd
 import cv2
-import time
 import os
 
 
@@ -60,20 +59,26 @@ def build_option(args):
 model = fd.vision.tracking.PPTracking(
     model_file, params_file, config_file, runtime_option=runtime_option)
 
+# 初始化轨迹记录器
+recorder = fd.vision.tracking.TrailRecorder()
+# 绑定记录器 注意：每次预测时，往trail_recorder里面插入数据，随着预测次数的增加，内存会不断地增长，
+# 可以通过unbind_recorder()方法来解除绑定
+model.bind_recorder(recorder)
 # 预测图片分割结果
 cap = cv2.VideoCapture(args.video)
-frame_id = 0
+# count = 0
 while True:
-    start_time = time.time()
-    frame_id = frame_id+1
     _, frame = cap.read()
     if frame is None:
         break
     result = model.predict(frame)
-    end_time = time.time()
-    fps = 1.0/(end_time-start_time)
-    img = fd.vision.vis_mot(frame, result, fps, frame_id)
+    # count += 1
+    # if count == 10:
+    #     model.unbind_recorder()
+    img = fd.vision.vis_mot(frame, result, 0.0, recorder)
     cv2.imshow("video", img)
-    cv2.waitKey(30)
+    if cv2.waitKey(30) == ord("q"):
+        break
+model.unbind_recorder()
 cap.release()
 cv2.destroyAllWindows()
diff --git a/fastdeploy/vision/common/result.h b/fastdeploy/vision/common/result.h
@@ -14,6 +14,7 @@
 #pragma once
 #include "fastdeploy/fastdeploy_model.h"
 #include "opencv2/core/core.hpp"
+#include <set>
 
 namespace fastdeploy {
 /** \brief All C++ FastDeploy Vision Models APIs are defined inside this namespace
@@ -171,6 +172,7 @@ struct FASTDEPLOY_DECL MOTResult : public BaseResult {
   /** \brief The classify label id for all the tracking object
    */
   std::vector<int> class_ids;
+
   ResultType type = ResultType::MOT;
   /// Clear MOT result
   void Clear();

diff --git a/fastdeploy/vision/tracking/pptracking/model.cc b/fastdeploy/vision/tracking/pptracking/model.cc
@@ -161,9 +161,7 @@ bool PPTracking::Initialize() {
     return false;
   }
   // create JDETracker instance
-  std::unique_ptr<JDETracker> jdeTracker(new JDETracker);
-  jdeTracker_ = std::move(jdeTracker);
-
+  jdeTracker_ = std::unique_ptr<JDETracker>(new JDETracker);
   return true;
 }
 
@@ -245,7 +243,6 @@ bool PPTracking::Postprocess(std::vector<FDTensor>& infer_result, MOTResult *res
   cv::Mat dets(bbox_shape[0], 6, CV_32FC1, bbox_data);
   cv::Mat emb(bbox_shape[0], emb_shape[1], CV_32FC1, emb_data);
 
-
   result->Clear();
   std::vector<Track> tracks;
   std::vector<int> valid;
@@ -264,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) {
@@ -285,9 +281,36 @@ bool PPTracking::Postprocess(std::vector<FDTensor>& infer_result, MOTResult *res
       }
     }
   }
+  if (!is_record_trail_) return true;
+  int nums = result->boxes.size();
+  for (int i=0; i<nums; i++) {
+    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];
+    recorder_->Add(id,{int(center_x), int(center_y)});
+  }
   return true;
 }
 
+void PPTracking::BindRecorder(TrailRecorder* recorder){
+
+    recorder_ = recorder;
+    is_record_trail_ = true;
+}
+
+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++){
+      iter->second.clear();
+      iter->second.shrink_to_fit();
+    }
+    recorder_->records.clear();
+    std::map<int, std::vector<std::array<int, 2>>>().swap(recorder_->records);
+    recorder_ = nullptr;
+}
+
 } // namespace tracking
 } // namespace vision
 } // namespace fastdeploy
diff --git a/fastdeploy/vision/tracking/pptracking/model.h b/fastdeploy/vision/tracking/pptracking/model.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <map>
 #include "fastdeploy/vision/common/processors/transform.h"
 #include "fastdeploy/fastdeploy_model.h"
 #include "fastdeploy/vision/common/result.h"
@@ -22,6 +23,21 @@
 namespace fastdeploy {
 namespace vision {
 namespace tracking {
+struct TrailRecorder{
+  std::map<int, std::vector<std::array<int, 2>>> records;
+  void Add(int id, const std::array<int, 2>& 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 {
  public:
@@ -49,6 +65,14 @@ class FASTDEPLOY_DECL PPTracking: public FastDeployModel {
    * \return true if the prediction successed, otherwise false
    */
   virtual bool Predict(cv::Mat* img, MOTResult* result);
+  /** \brief bind tracking trail struct
+   *
+   * \param[in] recorder The MOT trail will record the trail of object
+   */
+  void BindRecorder(TrailRecorder* recorder);
+  /** \brief cancel binding and clear trail information
+   */
+  void UnbindRecorder();
 
  private:
   bool BuildPreprocessPipelineFromConfig();
@@ -65,8 +89,11 @@ class FASTDEPLOY_DECL PPTracking: public FastDeployModel {
   float conf_thresh_;
   float tracked_thresh_;
   float min_box_area_;
+  bool is_record_trail_ = false;
   std::unique_ptr<JDETracker> jdeTracker_;
+  TrailRecorder *recorder_ = nullptr;
 };
+
 }  // namespace tracking
 }  // namespace vision
 }  // namespace fastdeploy
diff --git a/fastdeploy/vision/tracking/pptracking/pptracking_pybind.cc b/fastdeploy/vision/tracking/pptracking/pptracking_pybind.cc
@@ -15,6 +15,11 @@
 
 namespace fastdeploy {
 void BindPPTracking(pybind11::module &m) {
+
+  pybind11::class_<vision::tracking::TrailRecorder>(m, "TrailRecorder")
+    .def(pybind11::init<>())
+    .def_readwrite("records", &vision::tracking::TrailRecorder::records)
+    .def("add", &vision::tracking::TrailRecorder::Add);
   pybind11::class_<vision::tracking::PPTracking, FastDeployModel>(
     m, "PPTracking")
     .def(pybind11::init<std::string, std::string, std::string, RuntimeOption,
@@ -26,6 +31,8 @@ void BindPPTracking(pybind11::module &m) {
              vision::MOTResult *res = new vision::MOTResult();
              self.Predict(&mat, res);
              return res;
-         });
+         })
+    .def("bind_recorder", &vision::tracking::PPTracking::BindRecorder)
+    .def("unbind_recorder", &vision::tracking::PPTracking::UnbindRecorder);
 }
 }  // namespace fastdeploy
diff --git a/fastdeploy/vision/tracking/pptracking/trajectory.cc b/fastdeploy/vision/tracking/pptracking/trajectory.cc
@@ -118,7 +118,7 @@ void Trajectory::update(Trajectory *traj,
   if (update_embedding_) update_embedding(traj->current_embedding);
 }
 
-void Trajectory::activate(int &cnt,int timestamp_) {
+void Trajectory::activate(int &cnt, int timestamp_) {
   id = next_id(cnt);
   TKalmanFilter::init(cv::Mat(xyah));
   length = 0;
@@ -130,7 +130,7 @@ void Trajectory::activate(int &cnt,int timestamp_) {
   starttime = timestamp_;
 }
 
-void Trajectory::reactivate(Trajectory *traj,int &cnt, int timestamp_, bool newid) {
+void Trajectory::reactivate(Trajectory *traj, int &cnt, int timestamp_, bool newid) {
   TKalmanFilter::correct(cv::Mat(traj->xyah));
   update_embedding(traj->current_embedding);
   length = 0;

diff --git a/fastdeploy/vision/tracking/pptracking/trajectory.h b/fastdeploy/vision/tracking/pptracking/trajectory.h
@@ -74,8 +74,8 @@ class FASTDEPLOY_DECL Trajectory : public TKalmanFilter {
   virtual void update(Trajectory *traj,
                       int timestamp,
                       bool update_embedding = true);
-  virtual void activate(int& cnt, int timestamp);
-  virtual void reactivate(Trajectory *traj, int & cnt,int timestamp, bool newid = false);
+  virtual void activate(int &cnt, int timestamp);
+  virtual void reactivate(Trajectory *traj, int &cnt, int timestamp, bool newid = false);
   virtual void mark_lost(void);
   virtual void mark_removed(void);
Original file line number	Diff line number	Diff line change
Expand Up		@@ -37,4 +37,3 @@ fastdeploy.vision.MOTResult
		- ids(list of list(float)):成员变量，表示单帧画面中所有目标的id，其元素个数与`boxes`一致
		- scores(list of float): 成员变量，表示单帧画面检测出来的所有目标置信度
		- class_ids(list of int): 成员变量，表示单帧画面出来的所有目标类别