更换文档检测模型

paddle_detection/deploy/pptracking/cpp/include/pipeline.h

@@ -0,0 +1,142 @@
+//   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DEPLOY_PPTRACKING_CPP_INCLUDE_PIPELINE_H_
+#define DEPLOY_PPTRACKING_CPP_INCLUDE_PIPELINE_H_
+
+#include <glog/logging.h>
+
+#include <math.h>
+#include <sys/types.h>
+#include <algorithm>
+#include <iostream>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#ifdef _WIN32
+#include <direct.h>
+#include <io.h>
+#elif LINUX
+#include <stdarg.h>
+#include <sys/stat.h>
+#endif
+
+#include "include/jde_predictor.h"
+#include "include/sde_predictor.h"
+
+namespace PaddleDetection {
+
+class Pipeline {
+ public:
+  explicit Pipeline(const std::string& device,
+                    const double threshold,
+                    const std::string& output_dir,
+                    const std::string& run_mode = "paddle",
+                    const int gpu_id = 0,
+                    const bool use_mkldnn = false,
+                    const int cpu_threads = 1,
+                    const bool trt_calib_mode = false,
+                    const bool do_entrance_counting = false,
+                    const bool save_result = false,
+                    const std::string& scene = "pedestrian",
+                    const bool tiny_obj = false,
+                    const bool is_mtmct = false,
+                    const int secs_interval = 10,
+                    const std::string track_model_dir = "",
+                    const std::string det_model_dir = "",
+                    const std::string reid_model_dir = "") {
+    std::vector<std::string> input;
+    this->input_ = input;
+    this->device_ = device;
+    this->threshold_ = threshold;
+    this->output_dir_ = output_dir;
+    this->run_mode_ = run_mode;
+    this->gpu_id_ = gpu_id;
+    this->use_mkldnn_ = use_mkldnn;
+    this->cpu_threads_ = cpu_threads;
+    this->trt_calib_mode_ = trt_calib_mode;
+    this->do_entrance_counting_ = do_entrance_counting;
+    this->secs_interval_ = secs_interval_;
+    this->save_result_ = save_result;
+    SelectModel(scene,
+                tiny_obj,
+                is_mtmct,
+                track_model_dir,
+                det_model_dir,
+                reid_model_dir);
+    InitPredictor();
+  }
+
+  // Set input, it must execute before Run()
+  void SetInput(const std::string& input_video);
+  void ClearInput();
+
+  // Run pipeline in video
+  void Run();
+  void PredictMOT(const std::string& video_path);
+  void PredictMTMCT(const std::vector<std::string> video_inputs);
+
+  // Run pipeline in stream
+  void RunMOTStream(const cv::Mat img,
+                    const int frame_id,
+                    const int video_fps,
+                    const Rect entrance,
+                    cv::Mat out_img,
+                    std::vector<std::string>* records,
+                    std::set<int>* count_set,
+                    std::set<int>* interval_count_set,
+                    std::vector<int>* in_count_list,
+                    std::vector<int>* out_count_list,
+                    std::map<int, std::vector<float>>* prev_center,
+                    std::vector<std::string>* flow_records);
+  void RunMTMCTStream(const std::vector<cv::Mat> imgs,
+                      std::vector<std::string>* records);
+
+  void PrintBenchmarkLog(const std::vector<double> det_time, const int img_num);
+
+ private:
+  // Select model according to scenes, it must execute before Run()
+  void SelectModel(const std::string& scene = "pedestrian",
+                   const bool tiny_obj = false,
+                   const bool is_mtmct = false,
+                   const std::string track_model_dir = "",
+                   const std::string det_model_dir = "",
+                   const std::string reid_model_dir = "");
+  void InitPredictor();
+
+  std::shared_ptr<PaddleDetection::JDEPredictor> jde_sct_;
+  std::shared_ptr<PaddleDetection::SDEPredictor> sde_sct_;
+
+  std::vector<std::string> input_;
+  std::vector<cv::Mat> stream_;
+  std::string device_;
+  double threshold_;
+  std::string output_dir_;
+  std::string track_model_dir_;
+  std::string det_model_dir_;
+  std::string reid_model_dir_;
+  std::string run_mode_ = "paddle";
+  int gpu_id_ = 0;
+  bool use_mkldnn_ = false;
+  int cpu_threads_ = 1;
+  bool trt_calib_mode_ = false;
+  bool do_entrance_counting_ = false;
+  bool save_result_ = false;
+  int secs_interval_ = 10;
+};
+
+}  // namespace PaddleDetection
+
+#endif  // DEPLOY_PPTRACKING_CPP_INCLUDE_PIPELINE_H_