ocr_db_crnn.cc 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680
  1. // Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #include <chrono>
  15. #include "paddle_api.h" // NOLINT
  16. #include "paddle_place.h"
  17. #include "cls_process.h"
  18. #include "crnn_process.h"
  19. #include "db_post_process.h"
  20. #include "AutoLog/auto_log/lite_autolog.h"
  21. using namespace paddle::lite_api; // NOLINT
  22. using namespace std;
  23. // fill tensor with mean and scale and trans layout: nhwc -> nchw, neon speed up
  24. void NeonMeanScale(const float *din, float *dout, int size,
  25. const std::vector<float> mean,
  26. const std::vector<float> scale) {
  27. if (mean.size() != 3 || scale.size() != 3) {
  28. std::cerr << "[ERROR] mean or scale size must equal to 3" << std::endl;
  29. exit(1);
  30. }
  31. float32x4_t vmean0 = vdupq_n_f32(mean[0]);
  32. float32x4_t vmean1 = vdupq_n_f32(mean[1]);
  33. float32x4_t vmean2 = vdupq_n_f32(mean[2]);
  34. float32x4_t vscale0 = vdupq_n_f32(scale[0]);
  35. float32x4_t vscale1 = vdupq_n_f32(scale[1]);
  36. float32x4_t vscale2 = vdupq_n_f32(scale[2]);
  37. float *dout_c0 = dout;
  38. float *dout_c1 = dout + size;
  39. float *dout_c2 = dout + size * 2;
  40. int i = 0;
  41. for (; i < size - 3; i += 4) {
  42. float32x4x3_t vin3 = vld3q_f32(din);
  43. float32x4_t vsub0 = vsubq_f32(vin3.val[0], vmean0);
  44. float32x4_t vsub1 = vsubq_f32(vin3.val[1], vmean1);
  45. float32x4_t vsub2 = vsubq_f32(vin3.val[2], vmean2);
  46. float32x4_t vs0 = vmulq_f32(vsub0, vscale0);
  47. float32x4_t vs1 = vmulq_f32(vsub1, vscale1);
  48. float32x4_t vs2 = vmulq_f32(vsub2, vscale2);
  49. vst1q_f32(dout_c0, vs0);
  50. vst1q_f32(dout_c1, vs1);
  51. vst1q_f32(dout_c2, vs2);
  52. din += 12;
  53. dout_c0 += 4;
  54. dout_c1 += 4;
  55. dout_c2 += 4;
  56. }
  57. for (; i < size; i++) {
  58. *(dout_c0++) = (*(din++) - mean[0]) * scale[0];
  59. *(dout_c1++) = (*(din++) - mean[1]) * scale[1];
  60. *(dout_c2++) = (*(din++) - mean[2]) * scale[2];
  61. }
  62. }
  63. // resize image to a size multiple of 32 which is required by the network
  64. cv::Mat DetResizeImg(const cv::Mat img, int max_size_len,
  65. std::vector<float> &ratio_hw) {
  66. int w = img.cols;
  67. int h = img.rows;
  68. float ratio = 1.f;
  69. int max_wh = w >= h ? w : h;
  70. if (max_wh > max_size_len) {
  71. if (h > w) {
  72. ratio = static_cast<float>(max_size_len) / static_cast<float>(h);
  73. } else {
  74. ratio = static_cast<float>(max_size_len) / static_cast<float>(w);
  75. }
  76. }
  77. int resize_h = static_cast<int>(float(h) * ratio);
  78. int resize_w = static_cast<int>(float(w) * ratio);
  79. if (resize_h % 32 == 0)
  80. resize_h = resize_h;
  81. else if (resize_h / 32 < 1 + 1e-5)
  82. resize_h = 32;
  83. else
  84. resize_h = (resize_h / 32 - 1) * 32;
  85. if (resize_w % 32 == 0)
  86. resize_w = resize_w;
  87. else if (resize_w / 32 < 1 + 1e-5)
  88. resize_w = 32;
  89. else
  90. resize_w = (resize_w / 32 - 1) * 32;
  91. cv::Mat resize_img;
  92. cv::resize(img, resize_img, cv::Size(resize_w, resize_h));
  93. ratio_hw.push_back(static_cast<float>(resize_h) / static_cast<float>(h));
  94. ratio_hw.push_back(static_cast<float>(resize_w) / static_cast<float>(w));
  95. return resize_img;
  96. }
  97. cv::Mat RunClsModel(cv::Mat img, std::shared_ptr<PaddlePredictor> predictor_cls,
  98. const float thresh = 0.9) {
  99. std::vector<float> mean = {0.5f, 0.5f, 0.5f};
  100. std::vector<float> scale = {1 / 0.5f, 1 / 0.5f, 1 / 0.5f};
  101. cv::Mat srcimg;
  102. img.copyTo(srcimg);
  103. cv::Mat crop_img;
  104. img.copyTo(crop_img);
  105. cv::Mat resize_img;
  106. int index = 0;
  107. float wh_ratio =
  108. static_cast<float>(crop_img.cols) / static_cast<float>(crop_img.rows);
  109. resize_img = ClsResizeImg(crop_img);
  110. resize_img.convertTo(resize_img, CV_32FC3, 1 / 255.f);
  111. const float *dimg = reinterpret_cast<const float *>(resize_img.data);
  112. std::unique_ptr<Tensor> input_tensor0(std::move(predictor_cls->GetInput(0)));
  113. input_tensor0->Resize({1, 3, resize_img.rows, resize_img.cols});
  114. auto *data0 = input_tensor0->mutable_data<float>();
  115. NeonMeanScale(dimg, data0, resize_img.rows * resize_img.cols, mean, scale);
  116. // Run CLS predictor
  117. predictor_cls->Run();
  118. // Get output and run postprocess
  119. std::unique_ptr<const Tensor> softmax_out(
  120. std::move(predictor_cls->GetOutput(0)));
  121. auto *softmax_scores = softmax_out->mutable_data<float>();
  122. auto softmax_out_shape = softmax_out->shape();
  123. float score = 0;
  124. int label = 0;
  125. for (int i = 0; i < softmax_out_shape[1]; i++) {
  126. if (softmax_scores[i] > score) {
  127. score = softmax_scores[i];
  128. label = i;
  129. }
  130. }
  131. if (label % 2 == 1 && score > thresh) {
  132. cv::rotate(srcimg, srcimg, 1);
  133. }
  134. return srcimg;
  135. }
  136. void RunRecModel(std::vector<std::vector<std::vector<int>>> boxes, cv::Mat img,
  137. std::shared_ptr<PaddlePredictor> predictor_crnn,
  138. std::vector<std::string> &rec_text,
  139. std::vector<float> &rec_text_score,
  140. std::vector<std::string> charactor_dict,
  141. std::shared_ptr<PaddlePredictor> predictor_cls,
  142. int use_direction_classify,
  143. std::vector<double> *times,
  144. int rec_image_height) {
  145. std::vector<float> mean = {0.5f, 0.5f, 0.5f};
  146. std::vector<float> scale = {1 / 0.5f, 1 / 0.5f, 1 / 0.5f};
  147. cv::Mat srcimg;
  148. img.copyTo(srcimg);
  149. cv::Mat crop_img;
  150. cv::Mat resize_img;
  151. int index = 0;
  152. std::vector<double> time_info = {0, 0, 0};
  153. for (int i = boxes.size() - 1; i >= 0; i--) {
  154. auto preprocess_start = std::chrono::steady_clock::now();
  155. crop_img = GetRotateCropImage(srcimg, boxes[i]);
  156. if (use_direction_classify >= 1) {
  157. crop_img = RunClsModel(crop_img, predictor_cls);
  158. }
  159. float wh_ratio =
  160. static_cast<float>(crop_img.cols) / static_cast<float>(crop_img.rows);
  161. resize_img = CrnnResizeImg(crop_img, wh_ratio, rec_image_height);
  162. resize_img.convertTo(resize_img, CV_32FC3, 1 / 255.f);
  163. const float *dimg = reinterpret_cast<const float *>(resize_img.data);
  164. std::unique_ptr<Tensor> input_tensor0(
  165. std::move(predictor_crnn->GetInput(0)));
  166. input_tensor0->Resize({1, 3, resize_img.rows, resize_img.cols});
  167. auto *data0 = input_tensor0->mutable_data<float>();
  168. NeonMeanScale(dimg, data0, resize_img.rows * resize_img.cols, mean, scale);
  169. auto preprocess_end = std::chrono::steady_clock::now();
  170. //// Run CRNN predictor
  171. auto inference_start = std::chrono::steady_clock::now();
  172. predictor_crnn->Run();
  173. // Get output and run postprocess
  174. std::unique_ptr<const Tensor> output_tensor0(
  175. std::move(predictor_crnn->GetOutput(0)));
  176. auto *predict_batch = output_tensor0->data<float>();
  177. auto predict_shape = output_tensor0->shape();
  178. auto inference_end = std::chrono::steady_clock::now();
  179. // ctc decode
  180. auto postprocess_start = std::chrono::steady_clock::now();
  181. std::string str_res;
  182. int argmax_idx;
  183. int last_index = 0;
  184. float score = 0.f;
  185. int count = 0;
  186. float max_value = 0.0f;
  187. for (int n = 0; n < predict_shape[1]; n++) {
  188. argmax_idx = int(Argmax(&predict_batch[n * predict_shape[2]],
  189. &predict_batch[(n + 1) * predict_shape[2]]));
  190. max_value =
  191. float(*std::max_element(&predict_batch[n * predict_shape[2]],
  192. &predict_batch[(n + 1) * predict_shape[2]]));
  193. if (argmax_idx > 0 && (!(n > 0 && argmax_idx == last_index))) {
  194. score += max_value;
  195. count += 1;
  196. str_res += charactor_dict[argmax_idx];
  197. }
  198. last_index = argmax_idx;
  199. }
  200. score /= count;
  201. rec_text.push_back(str_res);
  202. rec_text_score.push_back(score);
  203. auto postprocess_end = std::chrono::steady_clock::now();
  204. std::chrono::duration<float> preprocess_diff = preprocess_end - preprocess_start;
  205. time_info[0] += double(preprocess_diff.count() * 1000);
  206. std::chrono::duration<float> inference_diff = inference_end - inference_start;
  207. time_info[1] += double(inference_diff.count() * 1000);
  208. std::chrono::duration<float> postprocess_diff = postprocess_end - postprocess_start;
  209. time_info[2] += double(postprocess_diff.count() * 1000);
  210. }
  211. times->push_back(time_info[0]);
  212. times->push_back(time_info[1]);
  213. times->push_back(time_info[2]);
  214. }
  215. std::vector<std::vector<std::vector<int>>>
  216. RunDetModel(std::shared_ptr<PaddlePredictor> predictor, cv::Mat img,
  217. std::map<std::string, double> Config, std::vector<double> *times) {
  218. // Read img
  219. int max_side_len = int(Config["max_side_len"]);
  220. int det_db_use_dilate = int(Config["det_db_use_dilate"]);
  221. cv::Mat srcimg;
  222. img.copyTo(srcimg);
  223. auto preprocess_start = std::chrono::steady_clock::now();
  224. std::vector<float> ratio_hw;
  225. img = DetResizeImg(img, max_side_len, ratio_hw);
  226. cv::Mat img_fp;
  227. img.convertTo(img_fp, CV_32FC3, 1.0 / 255.f);
  228. // Prepare input data from image
  229. std::unique_ptr<Tensor> input_tensor0(std::move(predictor->GetInput(0)));
  230. input_tensor0->Resize({1, 3, img_fp.rows, img_fp.cols});
  231. auto *data0 = input_tensor0->mutable_data<float>();
  232. std::vector<float> mean = {0.485f, 0.456f, 0.406f};
  233. std::vector<float> scale = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f};
  234. const float *dimg = reinterpret_cast<const float *>(img_fp.data);
  235. NeonMeanScale(dimg, data0, img_fp.rows * img_fp.cols, mean, scale);
  236. auto preprocess_end = std::chrono::steady_clock::now();
  237. // Run predictor
  238. auto inference_start = std::chrono::steady_clock::now();
  239. predictor->Run();
  240. // Get output and post process
  241. std::unique_ptr<const Tensor> output_tensor(
  242. std::move(predictor->GetOutput(0)));
  243. auto *outptr = output_tensor->data<float>();
  244. auto shape_out = output_tensor->shape();
  245. auto inference_end = std::chrono::steady_clock::now();
  246. // Save output
  247. auto postprocess_start = std::chrono::steady_clock::now();
  248. float pred[shape_out[2] * shape_out[3]];
  249. unsigned char cbuf[shape_out[2] * shape_out[3]];
  250. for (int i = 0; i < int(shape_out[2] * shape_out[3]); i++) {
  251. pred[i] = static_cast<float>(outptr[i]);
  252. cbuf[i] = static_cast<unsigned char>((outptr[i]) * 255);
  253. }
  254. cv::Mat cbuf_map(shape_out[2], shape_out[3], CV_8UC1,
  255. reinterpret_cast<unsigned char *>(cbuf));
  256. cv::Mat pred_map(shape_out[2], shape_out[3], CV_32F,
  257. reinterpret_cast<float *>(pred));
  258. const double threshold = double(Config["det_db_thresh"]) * 255;
  259. const double max_value = 255;
  260. cv::Mat bit_map;
  261. cv::threshold(cbuf_map, bit_map, threshold, max_value, cv::THRESH_BINARY);
  262. if (det_db_use_dilate == 1) {
  263. cv::Mat dilation_map;
  264. cv::Mat dila_ele =
  265. cv::getStructuringElement(cv::MORPH_RECT, cv::Size(2, 2));
  266. cv::dilate(bit_map, dilation_map, dila_ele);
  267. bit_map = dilation_map;
  268. }
  269. auto boxes = BoxesFromBitmap(pred_map, bit_map, Config);
  270. std::vector<std::vector<std::vector<int>>> filter_boxes =
  271. FilterTagDetRes(boxes, ratio_hw[0], ratio_hw[1], srcimg);
  272. auto postprocess_end = std::chrono::steady_clock::now();
  273. std::chrono::duration<float> preprocess_diff = preprocess_end - preprocess_start;
  274. times->push_back(double(preprocess_diff.count() * 1000));
  275. std::chrono::duration<float> inference_diff = inference_end - inference_start;
  276. times->push_back(double(inference_diff.count() * 1000));
  277. std::chrono::duration<float> postprocess_diff = postprocess_end - postprocess_start;
  278. times->push_back(double(postprocess_diff.count() * 1000));
  279. return filter_boxes;
  280. }
  281. std::shared_ptr<PaddlePredictor> loadModel(std::string model_file, int num_threads) {
  282. MobileConfig config;
  283. config.set_model_from_file(model_file);
  284. config.set_threads(num_threads);
  285. std::shared_ptr<PaddlePredictor> predictor =
  286. CreatePaddlePredictor<MobileConfig>(config);
  287. return predictor;
  288. }
  289. cv::Mat Visualization(cv::Mat srcimg,
  290. std::vector<std::vector<std::vector<int>>> boxes) {
  291. cv::Point rook_points[boxes.size()][4];
  292. for (int n = 0; n < boxes.size(); n++) {
  293. for (int m = 0; m < boxes[0].size(); m++) {
  294. rook_points[n][m] = cv::Point(static_cast<int>(boxes[n][m][0]),
  295. static_cast<int>(boxes[n][m][1]));
  296. }
  297. }
  298. cv::Mat img_vis;
  299. srcimg.copyTo(img_vis);
  300. for (int n = 0; n < boxes.size(); n++) {
  301. const cv::Point *ppt[1] = {rook_points[n]};
  302. int npt[] = {4};
  303. cv::polylines(img_vis, ppt, npt, 1, 1, CV_RGB(0, 255, 0), 2, 8, 0);
  304. }
  305. cv::imwrite("./vis.jpg", img_vis);
  306. std::cout << "The detection visualized image saved in ./vis.jpg" << std::endl;
  307. return img_vis;
  308. }
  309. std::vector<std::string> split(const std::string &str,
  310. const std::string &delim) {
  311. std::vector<std::string> res;
  312. if ("" == str)
  313. return res;
  314. char *strs = new char[str.length() + 1];
  315. std::strcpy(strs, str.c_str());
  316. char *d = new char[delim.length() + 1];
  317. std::strcpy(d, delim.c_str());
  318. char *p = std::strtok(strs, d);
  319. while (p) {
  320. string s = p;
  321. res.push_back(s);
  322. p = std::strtok(NULL, d);
  323. }
  324. return res;
  325. }
  326. std::map<std::string, double> LoadConfigTxt(std::string config_path) {
  327. auto config = ReadDict(config_path);
  328. std::map<std::string, double> dict;
  329. for (int i = 0; i < config.size(); i++) {
  330. std::vector<std::string> res = split(config[i], " ");
  331. dict[res[0]] = stod(res[1]);
  332. }
  333. return dict;
  334. }
  335. void check_params(int argc, char **argv) {
  336. if (argc<=1 || (strcmp(argv[1], "det")!=0 && strcmp(argv[1], "rec")!=0 && strcmp(argv[1], "system")!=0)) {
  337. std::cerr << "Please choose one mode of [det, rec, system] !" << std::endl;
  338. exit(1);
  339. }
  340. if (strcmp(argv[1], "det") == 0) {
  341. if (argc < 9){
  342. std::cerr << "[ERROR] usage:" << argv[0]
  343. << " det det_model runtime_device num_threads batchsize img_dir det_config lite_benchmark_value" << std::endl;
  344. exit(1);
  345. }
  346. }
  347. if (strcmp(argv[1], "rec") == 0) {
  348. if (argc < 9){
  349. std::cerr << "[ERROR] usage:" << argv[0]
  350. << " rec rec_model runtime_device num_threads batchsize img_dir key_txt lite_benchmark_value" << std::endl;
  351. exit(1);
  352. }
  353. }
  354. if (strcmp(argv[1], "system") == 0) {
  355. if (argc < 12){
  356. std::cerr << "[ERROR] usage:" << argv[0]
  357. << " system det_model rec_model clas_model runtime_device num_threads batchsize img_dir det_config key_txt lite_benchmark_value" << std::endl;
  358. exit(1);
  359. }
  360. }
  361. }
  362. void system(char **argv){
  363. std::string det_model_file = argv[2];
  364. std::string rec_model_file = argv[3];
  365. std::string cls_model_file = argv[4];
  366. std::string runtime_device = argv[5];
  367. std::string precision = argv[6];
  368. std::string num_threads = argv[7];
  369. std::string batchsize = argv[8];
  370. std::string img_dir = argv[9];
  371. std::string det_config_path = argv[10];
  372. std::string dict_path = argv[11];
  373. if (strcmp(argv[6], "FP32") != 0 && strcmp(argv[6], "INT8") != 0) {
  374. std::cerr << "Only support FP32 or INT8." << std::endl;
  375. exit(1);
  376. }
  377. std::vector<cv::String> cv_all_img_names;
  378. cv::glob(img_dir, cv_all_img_names);
  379. //// load config from txt file
  380. auto Config = LoadConfigTxt(det_config_path);
  381. int use_direction_classify = int(Config["use_direction_classify"]);
  382. int rec_image_height = int(Config["rec_image_height"]);
  383. auto charactor_dict = ReadDict(dict_path);
  384. charactor_dict.insert(charactor_dict.begin(), "#"); // blank char for ctc
  385. charactor_dict.push_back(" ");
  386. auto det_predictor = loadModel(det_model_file, std::stoi(num_threads));
  387. auto rec_predictor = loadModel(rec_model_file, std::stoi(num_threads));
  388. auto cls_predictor = loadModel(cls_model_file, std::stoi(num_threads));
  389. std::vector<double> det_time_info = {0, 0, 0};
  390. std::vector<double> rec_time_info = {0, 0, 0};
  391. for (int i = 0; i < cv_all_img_names.size(); ++i) {
  392. std::cout << "The predict img: " << cv_all_img_names[i] << std::endl;
  393. cv::Mat srcimg = cv::imread(cv_all_img_names[i], cv::IMREAD_COLOR);
  394. if (!srcimg.data) {
  395. std::cerr << "[ERROR] image read failed! image path: " << cv_all_img_names[i] << std::endl;
  396. exit(1);
  397. }
  398. std::vector<double> det_times;
  399. auto boxes = RunDetModel(det_predictor, srcimg, Config, &det_times);
  400. std::vector<std::string> rec_text;
  401. std::vector<float> rec_text_score;
  402. std::vector<double> rec_times;
  403. RunRecModel(boxes, srcimg, rec_predictor, rec_text, rec_text_score,
  404. charactor_dict, cls_predictor, use_direction_classify, &rec_times, rec_image_height);
  405. //// visualization
  406. auto img_vis = Visualization(srcimg, boxes);
  407. //// print recognized text
  408. for (int i = 0; i < rec_text.size(); i++) {
  409. std::cout << i << "\t" << rec_text[i] << "\t" << rec_text_score[i]
  410. << std::endl;
  411. }
  412. det_time_info[0] += det_times[0];
  413. det_time_info[1] += det_times[1];
  414. det_time_info[2] += det_times[2];
  415. rec_time_info[0] += rec_times[0];
  416. rec_time_info[1] += rec_times[1];
  417. rec_time_info[2] += rec_times[2];
  418. }
  419. if (strcmp(argv[12], "True") == 0) {
  420. AutoLogger autolog_det(det_model_file,
  421. runtime_device,
  422. std::stoi(num_threads),
  423. std::stoi(batchsize),
  424. "dynamic",
  425. precision,
  426. det_time_info,
  427. cv_all_img_names.size());
  428. AutoLogger autolog_rec(rec_model_file,
  429. runtime_device,
  430. std::stoi(num_threads),
  431. std::stoi(batchsize),
  432. "dynamic",
  433. precision,
  434. rec_time_info,
  435. cv_all_img_names.size());
  436. autolog_det.report();
  437. std::cout << std::endl;
  438. autolog_rec.report();
  439. }
  440. }
  441. void det(int argc, char **argv) {
  442. std::string det_model_file = argv[2];
  443. std::string runtime_device = argv[3];
  444. std::string precision = argv[4];
  445. std::string num_threads = argv[5];
  446. std::string batchsize = argv[6];
  447. std::string img_dir = argv[7];
  448. std::string det_config_path = argv[8];
  449. if (strcmp(argv[4], "FP32") != 0 && strcmp(argv[4], "INT8") != 0) {
  450. std::cerr << "Only support FP32 or INT8." << std::endl;
  451. exit(1);
  452. }
  453. std::vector<cv::String> cv_all_img_names;
  454. cv::glob(img_dir, cv_all_img_names);
  455. //// load config from txt file
  456. auto Config = LoadConfigTxt(det_config_path);
  457. auto det_predictor = loadModel(det_model_file, std::stoi(num_threads));
  458. std::vector<double> time_info = {0, 0, 0};
  459. for (int i = 0; i < cv_all_img_names.size(); ++i) {
  460. std::cout << "The predict img: " << cv_all_img_names[i] << std::endl;
  461. cv::Mat srcimg = cv::imread(cv_all_img_names[i], cv::IMREAD_COLOR);
  462. if (!srcimg.data) {
  463. std::cerr << "[ERROR] image read failed! image path: " << cv_all_img_names[i] << std::endl;
  464. exit(1);
  465. }
  466. std::vector<double> times;
  467. auto boxes = RunDetModel(det_predictor, srcimg, Config, &times);
  468. //// visualization
  469. auto img_vis = Visualization(srcimg, boxes);
  470. std::cout << boxes.size() << " bboxes have detected:" << std::endl;
  471. for (int i=0; i<boxes.size(); i++){
  472. std::cout << "The " << i << " box:" << std::endl;
  473. for (int j=0; j<4; j++){
  474. for (int k=0; k<2; k++){
  475. std::cout << boxes[i][j][k] << "\t";
  476. }
  477. }
  478. std::cout << std::endl;
  479. }
  480. time_info[0] += times[0];
  481. time_info[1] += times[1];
  482. time_info[2] += times[2];
  483. }
  484. if (strcmp(argv[9], "True") == 0) {
  485. AutoLogger autolog(det_model_file,
  486. runtime_device,
  487. std::stoi(num_threads),
  488. std::stoi(batchsize),
  489. "dynamic",
  490. precision,
  491. time_info,
  492. cv_all_img_names.size());
  493. autolog.report();
  494. }
  495. }
  496. void rec(int argc, char **argv) {
  497. std::string rec_model_file = argv[2];
  498. std::string runtime_device = argv[3];
  499. std::string precision = argv[4];
  500. std::string num_threads = argv[5];
  501. std::string batchsize = argv[6];
  502. std::string img_dir = argv[7];
  503. std::string dict_path = argv[8];
  504. std::string config_path = argv[9];
  505. if (strcmp(argv[4], "FP32") != 0 && strcmp(argv[4], "INT8") != 0) {
  506. std::cerr << "Only support FP32 or INT8." << std::endl;
  507. exit(1);
  508. }
  509. auto Config = LoadConfigTxt(config_path);
  510. int rec_image_height = int(Config["rec_image_height"]);
  511. std::vector<cv::String> cv_all_img_names;
  512. cv::glob(img_dir, cv_all_img_names);
  513. auto charactor_dict = ReadDict(dict_path);
  514. charactor_dict.insert(charactor_dict.begin(), "#"); // blank char for ctc
  515. charactor_dict.push_back(" ");
  516. auto rec_predictor = loadModel(rec_model_file, std::stoi(num_threads));
  517. std::shared_ptr<PaddlePredictor> cls_predictor;
  518. std::vector<double> time_info = {0, 0, 0};
  519. for (int i = 0; i < cv_all_img_names.size(); ++i) {
  520. std::cout << "The predict img: " << cv_all_img_names[i] << std::endl;
  521. cv::Mat srcimg = cv::imread(cv_all_img_names[i], cv::IMREAD_COLOR);
  522. if (!srcimg.data) {
  523. std::cerr << "[ERROR] image read failed! image path: " << cv_all_img_names[i] << std::endl;
  524. exit(1);
  525. }
  526. int width = srcimg.cols;
  527. int height = srcimg.rows;
  528. std::vector<int> upper_left = {0, 0};
  529. std::vector<int> upper_right = {width, 0};
  530. std::vector<int> lower_right = {width, height};
  531. std::vector<int> lower_left = {0, height};
  532. std::vector<std::vector<int>> box = {upper_left, upper_right, lower_right, lower_left};
  533. std::vector<std::vector<std::vector<int>>> boxes = {box};
  534. std::vector<std::string> rec_text;
  535. std::vector<float> rec_text_score;
  536. std::vector<double> times;
  537. RunRecModel(boxes, srcimg, rec_predictor, rec_text, rec_text_score,
  538. charactor_dict, cls_predictor, 0, &times, rec_image_height);
  539. //// print recognized text
  540. for (int i = 0; i < rec_text.size(); i++) {
  541. std::cout << i << "\t" << rec_text[i] << "\t" << rec_text_score[i]
  542. << std::endl;
  543. }
  544. time_info[0] += times[0];
  545. time_info[1] += times[1];
  546. time_info[2] += times[2];
  547. }
  548. // TODO: support autolog
  549. if (strcmp(argv[9], "True") == 0) {
  550. AutoLogger autolog(rec_model_file,
  551. runtime_device,
  552. std::stoi(num_threads),
  553. std::stoi(batchsize),
  554. "dynamic",
  555. precision,
  556. time_info,
  557. cv_all_img_names.size());
  558. autolog.report();
  559. }
  560. }
  561. int main(int argc, char **argv) {
  562. check_params(argc, argv);
  563. std::cout << "mode: " << argv[1] << endl;
  564. if (strcmp(argv[1], "system") == 0) {
  565. system(argv);
  566. }
  567. if (strcmp(argv[1], "det") == 0) {
  568. det(argc, argv);
  569. }
  570. if (strcmp(argv[1], "rec") == 0) {
  571. rec(argc, argv);
  572. }
  573. return 0;
  574. }