提交 69a7a140 编写于 作者: littletomatodonkey's avatar littletomatodonkey

fix visualization

上级 483486af
...@@ -52,6 +52,8 @@ public: ...@@ -52,6 +52,8 @@ public:
this->rec_model_dir.assign(config_map_["rec_model_dir"]); this->rec_model_dir.assign(config_map_["rec_model_dir"]);
this->char_list_file.assign(config_map_["char_list_file"]); this->char_list_file.assign(config_map_["char_list_file"]);
this->visualize = bool(stoi(config_map_["visualize"]));
} }
bool use_gpu = false; bool use_gpu = false;
...@@ -76,6 +78,8 @@ public: ...@@ -76,6 +78,8 @@ public:
std::string char_list_file; std::string char_list_file;
bool visualize = true;
void PrintConfigInfo(); void PrintConfigInfo();
private: private:
......
...@@ -42,7 +42,8 @@ public: ...@@ -42,7 +42,8 @@ public:
const int &max_side_len = 960, const int &max_side_len = 960,
const double &det_db_thresh = 0.3, const double &det_db_thresh = 0.3,
const double &det_db_box_thresh = 0.5, 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); LoadModel(model_dir);
this->use_gpu_ = use_gpu; this->use_gpu_ = use_gpu;
...@@ -55,6 +56,8 @@ public: ...@@ -55,6 +56,8 @@ public:
this->det_db_thresh_ = det_db_thresh; this->det_db_thresh_ = det_db_thresh;
this->det_db_box_thresh_ = det_db_box_thresh; this->det_db_box_thresh_ = det_db_box_thresh;
this->det_db_unclip_ratio_ = det_db_unclip_ratio; this->det_db_unclip_ratio_ = det_db_unclip_ratio;
this->visualize_ = visualize;
} }
// Load Paddle inference model // Load Paddle inference model
...@@ -77,6 +80,8 @@ private: ...@@ -77,6 +80,8 @@ private:
double det_db_box_thresh_ = 0.5; double det_db_box_thresh_ = 0.5;
double det_db_unclip_ratio_ = 2.0; double det_db_unclip_ratio_ = 2.0;
bool visualize_ = true;
std::vector<float> mean_ = {0.485f, 0.456f, 0.406f}; std::vector<float> mean_ = {0.485f, 0.456f, 0.406f};
std::vector<float> scale_ = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f}; std::vector<float> scale_ = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f};
bool is_scale_ = true; bool is_scale_ = true;
......
...@@ -26,12 +26,20 @@ ...@@ -26,12 +26,20 @@
#include <fstream> #include <fstream>
#include <numeric> #include <numeric>
#include "opencv2/core.hpp"
#include "opencv2/imgcodecs.hpp"
#include "opencv2/imgproc.hpp"
namespace PaddleOCR { namespace PaddleOCR {
class Utility { class Utility {
public: public:
static std::vector<std::string> ReadDict(const std::string &path); 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> template <class ForwardIterator>
inline static size_t argmax(ForwardIterator first, ForwardIterator last) { inline static size_t argmax(ForwardIterator first, ForwardIterator last) {
return std::distance(first, std::max_element(first, last)); return std::distance(first, std::max_element(first, last));
......
...@@ -53,7 +53,8 @@ int main(int argc, char **argv) { ...@@ -53,7 +53,8 @@ int main(int argc, char **argv) {
DBDetector det(config.det_model_dir, config.use_gpu, config.gpu_id, DBDetector det(config.det_model_dir, config.use_gpu, config.gpu_id,
config.gpu_mem, config.cpu_math_library_num_threads, config.gpu_mem, config.cpu_math_library_num_threads,
config.max_side_len, config.det_db_thresh, 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, CRNNRecognizer rec(config.rec_model_dir, config.use_gpu, config.gpu_id,
config.gpu_mem, config.cpu_math_library_num_threads, config.gpu_mem, config.cpu_math_library_num_threads,
config.char_list_file); config.char_list_file);
......
...@@ -114,25 +114,9 @@ void DBDetector::Run(cv::Mat &img, ...@@ -114,25 +114,9 @@ void DBDetector::Run(cv::Mat &img,
boxes = post_processor_.FilterTagDetRes(boxes, ratio_h, ratio_w, srcimg); boxes = post_processor_.FilterTagDetRes(boxes, ratio_h, ratio_w, srcimg);
//// visualization //// visualization
cv::Point rook_points[boxes.size()][4]; if (this->visualize_) {
for (int n = 0; n < boxes.size(); n++) { Utility::VisualizeBboxes(srcimg, boxes);
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);
}
imwrite("./det_res.png", img_vis);
std::cout << "The detection visualized image saved in ./det_res.png"
<< std::endl;
} }
} // namespace PaddleOCR } // namespace PaddleOCR
\ No newline at end of file
...@@ -36,4 +36,26 @@ std::vector<std::string> Utility::ReadDict(const std::string &path) { ...@@ -36,4 +36,26 @@ std::vector<std::string> Utility::ReadDict(const std::string &path) {
return m_vec; 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 } // namespace PaddleOCR
\ No newline at end of file
...@@ -15,3 +15,7 @@ det_model_dir ./inference/det_db ...@@ -15,3 +15,7 @@ det_model_dir ./inference/det_db
rec_model_dir ./inference/rec_crnn rec_model_dir ./inference/rec_crnn
char_list_file ./tools/ppocr_keys_v1.txt char_list_file ./tools/ppocr_keys_v1.txt
img_path ../../doc/imgs/11.jpg img_path ../../doc/imgs/11.jpg
# show the detection results
visualize 1
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册