ocr_rec.h 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  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 "opencv2/core.hpp"
  15. #include "opencv2/imgcodecs.hpp"
  16. #include "opencv2/imgproc.hpp"
  17. #include "paddle_api.h"
  18. #include "paddle_inference_api.h"
  19. #include <chrono>
  20. #include <iomanip>
  21. #include <iostream>
  22. #include <ostream>
  23. #include <vector>
  24. #include <cstring>
  25. #include <fstream>
  26. #include <numeric>
  27. #include <include/ocr_cls.h>
  28. #include <include/postprocess_op.h>
  29. #include <include/preprocess_op.h>
  30. #include <include/utility.h>
  31. using namespace paddle_infer;
  32. namespace PaddleOCR {
  33. class CRNNRecognizer {
  34. public:
  35. explicit CRNNRecognizer(const std::string &model_dir, const bool &use_gpu,
  36. const int &gpu_id, const int &gpu_mem,
  37. const int &cpu_math_library_num_threads,
  38. const bool &use_mkldnn, const string &label_path,
  39. const bool &use_tensorrt, const bool &use_fp16) {
  40. this->use_gpu_ = use_gpu;
  41. this->gpu_id_ = gpu_id;
  42. this->gpu_mem_ = gpu_mem;
  43. this->cpu_math_library_num_threads_ = cpu_math_library_num_threads;
  44. this->use_mkldnn_ = use_mkldnn;
  45. this->use_tensorrt_ = use_tensorrt;
  46. this->use_fp16_ = use_fp16;
  47. this->label_list_ = Utility::ReadDict(label_path);
  48. this->label_list_.insert(this->label_list_.begin(),
  49. "#"); // blank char for ctc
  50. this->label_list_.push_back(" ");
  51. LoadModel(model_dir);
  52. }
  53. // Load Paddle inference model
  54. void LoadModel(const std::string &model_dir);
  55. void Run(std::vector<std::vector<std::vector<int>>> boxes, cv::Mat &img,
  56. Classifier *cls);
  57. private:
  58. std::shared_ptr<Predictor> predictor_;
  59. bool use_gpu_ = false;
  60. int gpu_id_ = 0;
  61. int gpu_mem_ = 4000;
  62. int cpu_math_library_num_threads_ = 4;
  63. bool use_mkldnn_ = false;
  64. std::vector<std::string> label_list_;
  65. std::vector<float> mean_ = {0.5f, 0.5f, 0.5f};
  66. std::vector<float> scale_ = {1 / 0.5f, 1 / 0.5f, 1 / 0.5f};
  67. bool is_scale_ = true;
  68. bool use_tensorrt_ = false;
  69. bool use_fp16_ = false;
  70. // pre-process
  71. CrnnResizeImg resize_op_;
  72. Normalize normalize_op_;
  73. Permute permute_op_;
  74. // post-process
  75. PostProcessor post_processor_;
  76. cv::Mat GetRotateCropImage(const cv::Mat &srcimage,
  77. std::vector<std::vector<int>> box);
  78. }; // class CrnnRecognizer
  79. } // namespace PaddleOCR