From 69a7a140cf5a0f423723063a67ad7e7d64861acf Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Mon, 13 Jul 2020 09:15:54 +0000 Subject: [PATCH] fix visualization --- deploy/cpp_infer/include/config.h | 4 ++++ deploy/cpp_infer/include/ocr_det.h | 7 ++++++- deploy/cpp_infer/include/utility.h | 8 ++++++++ deploy/cpp_infer/src/main.cpp | 3 ++- deploy/cpp_infer/src/ocr_det.cpp | 20 ++------------------ deploy/cpp_infer/src/utility.cpp | 22 ++++++++++++++++++++++ deploy/cpp_infer/tools/config.txt | 4 ++++ 7 files changed, 48 insertions(+), 20 deletions(-) diff --git a/deploy/cpp_infer/include/config.h b/deploy/cpp_infer/include/config.h index 37c5078a..4004fc04 100644 --- a/deploy/cpp_infer/include/config.h +++ b/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: diff --git a/deploy/cpp_infer/include/ocr_det.h b/deploy/cpp_infer/include/ocr_det.h index 3208d346..503d3ee3 100644 --- a/deploy/cpp_infer/include/ocr_det.h +++ b/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 mean_ = {0.485f, 0.456f, 0.406f}; std::vector scale_ = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f}; bool is_scale_ = true; diff --git a/deploy/cpp_infer/include/utility.h b/deploy/cpp_infer/include/utility.h index ebcc3e84..367e37e4 100644 --- a/deploy/cpp_infer/include/utility.h +++ b/deploy/cpp_infer/include/utility.h @@ -26,12 +26,20 @@ #include #include +#include "opencv2/core.hpp" +#include "opencv2/imgcodecs.hpp" +#include "opencv2/imgproc.hpp" + namespace PaddleOCR { class Utility { public: static std::vector ReadDict(const std::string &path); + static void + VisualizeBboxes(const cv::Mat &srcimg, + const std::vector>> &boxes); + template inline static size_t argmax(ForwardIterator first, ForwardIterator last) { return std::distance(first, std::max_element(first, last)); diff --git a/deploy/cpp_infer/src/main.cpp b/deploy/cpp_infer/src/main.cpp index b63bdafc..d5e3d4c6 100644 --- a/deploy/cpp_infer/src/main.cpp +++ b/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); diff --git a/deploy/cpp_infer/src/ocr_det.cpp b/deploy/cpp_infer/src/ocr_det.cpp index a449e1b3..56447b07 100644 --- a/deploy/cpp_infer/src/ocr_det.cpp +++ b/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 diff --git a/deploy/cpp_infer/src/utility.cpp b/deploy/cpp_infer/src/utility.cpp index 8ca5ac67..ffb74c2e 100644 --- a/deploy/cpp_infer/src/utility.cpp +++ b/deploy/cpp_infer/src/utility.cpp @@ -36,4 +36,26 @@ std::vector Utility::ReadDict(const std::string &path) { return m_vec; } +void Utility::VisualizeBboxes( + const cv::Mat &srcimg, + const std::vector>> &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 diff --git a/deploy/cpp_infer/tools/config.txt b/deploy/cpp_infer/tools/config.txt index 1ce1568e..03661bbb 100644 --- a/deploy/cpp_infer/tools/config.txt +++ b/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 + -- GitLab