fix visualization

deploy/cpp_infer/include/config.h

@@ -52,6 +52,8 @@ public:
     this->rec_model_dir.assign(config_map_["rec_model_dir"]);
 
     this->char_list_file.assign(config_map_["char_list_file"]);
+
+    this->visualize = bool(stoi(config_map_["visualize"]));
   }
 
   bool use_gpu = false;
@@ -76,6 +78,8 @@ public:
 
   std::string char_list_file;
 
+  bool visualize = true;
+
   void PrintConfigInfo();
 
 private:

deploy/cpp_infer/include/ocr_det.h

@@ -42,7 +42,8 @@ public:
                       const int &max_side_len = 960,
                       const double &det_db_thresh = 0.3,
                       const double &det_db_box_thresh = 0.5,
-                      const double &det_db_unclip_ratio = 2.0) {
+                      const double &det_db_unclip_ratio = 2.0,
+                      const bool &visualize = true) {
     LoadModel(model_dir);
 
     this->use_gpu_ = use_gpu;
@@ -55,6 +56,8 @@ public:
     this->det_db_thresh_ = det_db_thresh;
     this->det_db_box_thresh_ = det_db_box_thresh;
     this->det_db_unclip_ratio_ = det_db_unclip_ratio;
+
+    this->visualize_ = visualize;
   }
 
   // Load Paddle inference model
@@ -77,6 +80,8 @@ private:
   double det_db_box_thresh_ = 0.5;
   double det_db_unclip_ratio_ = 2.0;
 
+  bool visualize_ = true;
+
   std::vector<float> mean_ = {0.485f, 0.456f, 0.406f};
   std::vector<float> scale_ = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f};
   bool is_scale_ = true;

deploy/cpp_infer/include/utility.h

@@ -26,12 +26,20 @@
 #include <fstream>
 #include <numeric>
 
+#include "opencv2/core.hpp"
+#include "opencv2/imgcodecs.hpp"
+#include "opencv2/imgproc.hpp"
+
 namespace PaddleOCR {
 
 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);
+
   template <class ForwardIterator>
   inline static size_t argmax(ForwardIterator first, ForwardIterator last) {
     return std::distance(first, std::max_element(first, last));

deploy/cpp_infer/src/main.cpp

@@ -53,7 +53,8 @@ int main(int argc, char **argv) {
   DBDetector det(config.det_model_dir, config.use_gpu, config.gpu_id,
                  config.gpu_mem, config.cpu_math_library_num_threads,
                  config.max_side_len, config.det_db_thresh,
-                 config.det_db_box_thresh, config.det_db_unclip_ratio);
+                 config.det_db_box_thresh, config.det_db_unclip_ratio,
+                 config.visualize);
   CRNNRecognizer rec(config.rec_model_dir, config.use_gpu, config.gpu_id,
                      config.gpu_mem, config.cpu_math_library_num_threads,
                      config.char_list_file);

deploy/cpp_infer/src/ocr_det.cpp

@@ -114,25 +114,9 @@ void DBDetector::Run(cv::Mat &img,
   boxes = post_processor_.FilterTagDetRes(boxes, ratio_h, ratio_w, srcimg);
 
   //// visualization
-  cv::Point rook_points[boxes.size()][4];
-  for (int n = 0; n < boxes.size(); n++) {
-    for (int m = 0; m < boxes[0].size(); m++) {
-      rook_points[n][m] = cv::Point(int(boxes[n][m][0]), int(boxes[n][m][1]));
-    }
+  if (this->visualize_) {
+    Utility::VisualizeBboxes(srcimg, boxes);
   }
-
-  cv::Mat img_vis;
-  srcimg.copyTo(img_vis);
-  for (int n = 0; n < boxes.size(); n++) {
-    const cv::Point *ppt[1] = {rook_points[n]};
-    int npt[] = {4};
-    cv::polylines(img_vis, ppt, npt, 1, 1, CV_RGB(0, 255, 0), 2, 8, 0);
-  }
-
-  imwrite("./det_res.png", img_vis);
-
-  std::cout << "The detection visualized image saved in ./det_res.png"
-            << std::endl;
 }
 
 } // namespace PaddleOCR
\ No newline at end of file

deploy/cpp_infer/src/utility.cpp

@@ -36,4 +36,26 @@ std::vector<std::string> Utility::ReadDict(const std::string &path) {
   return m_vec;
 }
 
+void Utility::VisualizeBboxes(
+    const cv::Mat &srcimg,
+    const std::vector<std::vector<std::vector<int>>> &boxes) {
+  cv::Point rook_points[boxes.size()][4];
+  for (int n = 0; n < boxes.size(); n++) {
+    for (int m = 0; m < boxes[0].size(); m++) {
+      rook_points[n][m] = cv::Point(int(boxes[n][m][0]), int(boxes[n][m][1]));
+    }
+  }
+  cv::Mat img_vis;
+  srcimg.copyTo(img_vis);
+  for (int n = 0; n < boxes.size(); n++) {
+    const cv::Point *ppt[1] = {rook_points[n]};
+    int npt[] = {4};
+    cv::polylines(img_vis, ppt, npt, 1, 1, CV_RGB(0, 255, 0), 2, 8, 0);
+  }
+
+  cv::imwrite("./ocr_vis.png", img_vis);
+  std::cout << "The detection visualized image saved in ./ocr_vis.png.pn"
+            << std::endl;
+}
+
 } // namespace PaddleOCR
\ No newline at end of file

deploy/cpp_infer/tools/config.txt

@@ -15,3 +15,7 @@ det_model_dir  ./inference/det_db
 rec_model_dir  ./inference/rec_crnn
 char_list_file ./tools/ppocr_keys_v1.txt
 img_path  ../../doc/imgs/11.jpg
+
+# show the detection results
+visualize 1
+