fix visualization
This commit is contained in:
parent
483486af58
commit
69a7a140cf
|
@ -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:
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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));
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
||||
|
||||
|
|
Loading…
Reference in New Issue