Refactor the code to keep the logic consistent with the whl package

deploy/cpp_infer/include/args.h

@@ -27,7 +27,7 @@ DECLARE_string(precision);
 DECLARE_bool(benchmark);
 DECLARE_string(output);
 DECLARE_string(image_dir);
-DECLARE_bool(visualize);
+DECLARE_string(type);
 // detection related
 DECLARE_string(det_model_dir);
 DECLARE_int32(max_side_len);
@@ -36,11 +36,17 @@ DECLARE_double(det_db_box_thresh);
 DECLARE_double(det_db_unclip_ratio);
 DECLARE_bool(use_dilation);
 DECLARE_string(det_db_score_mode);
+DECLARE_bool(visualize);
 // classification related
 DECLARE_bool(use_angle_cls);
 DECLARE_string(cls_model_dir);
 DECLARE_double(cls_thresh);
+DECLARE_int32(cls_batch_num);
 // recognition related
 DECLARE_string(rec_model_dir);
 DECLARE_int32(rec_batch_num);
-DECLARE_string(rec_char_dict_path);
+DECLARE_string(rec_char_dict_path);
+// forward related
+DECLARE_bool(det);
+DECLARE_bool(rec);
+DECLARE_bool(cls);

deploy/cpp_infer/include/ocr_cls.h

@@ -42,7 +42,8 @@ public:
                       const int &gpu_id, const int &gpu_mem,
                       const int &cpu_math_library_num_threads,
                       const bool &use_mkldnn, const double &cls_thresh,
-                      const bool &use_tensorrt, const std::string &precision) {
+                      const bool &use_tensorrt, const std::string &precision,
+                      const int &cls_batch_num) {
     this->use_gpu_ = use_gpu;
     this->gpu_id_ = gpu_id;
     this->gpu_mem_ = gpu_mem;
@@ -52,14 +53,17 @@ public:
     this->cls_thresh = cls_thresh;
     this->use_tensorrt_ = use_tensorrt;
     this->precision_ = precision;
+    this->cls_batch_num_ = cls_batch_num;
 
     LoadModel(model_dir);
   }
+  double cls_thresh = 0.5;
 
   // Load Paddle inference model
   void LoadModel(const std::string &model_dir);
 
-  cv::Mat Run(cv::Mat &img);
+  void Run(std::vector<cv::Mat> img_list, std::vector<int> &cls_labels,
+           std::vector<float> &cls_scores, std::vector<double> &times);
 
 private:
   std::shared_ptr<Predictor> predictor_;
@@ -69,17 +73,17 @@ private:
   int gpu_mem_ = 4000;
   int cpu_math_library_num_threads_ = 4;
   bool use_mkldnn_ = false;
-  double cls_thresh = 0.5;
 
   std::vector<float> mean_ = {0.5f, 0.5f, 0.5f};
   std::vector<float> scale_ = {1 / 0.5f, 1 / 0.5f, 1 / 0.5f};
   bool is_scale_ = true;
   bool use_tensorrt_ = false;
   std::string precision_ = "fp32";
+  int cls_batch_num_ = 1;
   // pre-process
   ClsResizeImg resize_op_;
   Normalize normalize_op_;
-  Permute permute_op_;
+  PermuteBatch permute_op_;
 
 }; // class Classifier
 

deploy/cpp_infer/include/ocr_det.h

@@ -73,7 +73,7 @@ public:
 
   // Run predictor
   void Run(cv::Mat &img, std::vector<std::vector<std::vector<int>>> &boxes,
-           std::vector<double> *times);
+           std::vector<double> &times);
 
 private:
   std::shared_ptr<Predictor> predictor_;

deploy/cpp_infer/include/ocr_rec.h

@@ -68,7 +68,7 @@ public:
   void LoadModel(const std::string &model_dir);
 
   void Run(std::vector<cv::Mat> img_list, std::vector<std::string> &rec_texts,
-           std::vector<float> &rec_text_scores, std::vector<double> *times);
+           std::vector<float> &rec_text_scores, std::vector<double> &times);
 
 private:
   std::shared_ptr<Predictor> predictor_;

deploy/cpp_infer/include/paddleocr.h

@@ -0,0 +1,67 @@
+// Copyright (c) 2020 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.
+
+#pragma once
+
+#include "opencv2/core.hpp"
+#include "opencv2/imgcodecs.hpp"
+#include "opencv2/imgproc.hpp"
+#include "paddle_api.h"
+#include "paddle_inference_api.h"
+#include <chrono>
+#include <iomanip>
+#include <iostream>
+#include <ostream>
+#include <vector>
+
+#include <cstring>
+#include <fstream>
+#include <numeric>
+
+#include <include/ocr_cls.h>
+#include <include/ocr_det.h>
+#include <include/ocr_rec.h>
+#include <include/preprocess_op.h>
+#include <include/utility.h>
+
+using namespace paddle_infer;
+
+namespace PaddleOCR {
+
+class PaddleOCR {
+public:
+  explicit PaddleOCR();
+  ~PaddleOCR();
+  std::vector<std::vector<OCRPredictResult>>
+  ocr(std::vector<cv::String> cv_all_img_names, bool det = true,
+      bool rec = true, bool cls = true);
+
+private:
+  DBDetector *detector_ = nullptr;
+  Classifier *classifier_ = nullptr;
+  CRNNRecognizer *recognizer_ = nullptr;
+
+  void det(cv::Mat img, std::vector<OCRPredictResult> &ocr_results,
+           std::vector<double> &times);
+  void rec(std::vector<cv::Mat> img_list,
+           std::vector<OCRPredictResult> &ocr_results,
+           std::vector<double> &times);
+  void cls(std::vector<cv::Mat> img_list,
+           std::vector<OCRPredictResult> &ocr_results,
+           std::vector<double> &times);
+  void log(std::vector<double> &det_times, std::vector<double> &rec_times,
+           std::vector<double> &cls_times, int img_num);
+};
+
+} // namespace PaddleOCR

deploy/cpp_infer/include/utility.h

@@ -32,14 +32,21 @@
 
 namespace PaddleOCR {
 
+struct OCRPredictResult {
+  std::vector<std::vector<int>> box;
+  std::string text;
+  float score = -1.0;
+  float cls_score;
+  int cls_label = -1;
+};
+
 class Utility {
 public:
   static std::vector<std::string> ReadDict(const std::string &path);
 
-  static void
-  VisualizeBboxes(const cv::Mat &srcimg,
-                  const std::vector<std::vector<std::vector<int>>> &boxes,
-                  const std::string &save_path);
+  static void VisualizeBboxes(const cv::Mat &srcimg,
+                              const std::vector<OCRPredictResult> &ocr_result,
+                              const std::string &save_path);
 
   template <class ForwardIterator>
   inline static size_t argmax(ForwardIterator first, ForwardIterator last) {
@@ -55,6 +62,10 @@ public:
   static std::vector<int> argsort(const std::vector<float> &array);
 
   static std::string basename(const std::string &filename);
+
+  static bool PathExists(const std::string &path);
+
+  static void print_result(const std::vector<OCRPredictResult> &ocr_result);
 };
 
 } // namespace PaddleOCR