config.h 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  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. #pragma once
  15. #include <iomanip>
  16. #include <iostream>
  17. #include <map>
  18. #include <ostream>
  19. #include <string>
  20. #include <vector>
  21. #include "include/utility.h"
  22. namespace PaddleOCR {
  23. class OCRConfig {
  24. public:
  25. explicit OCRConfig(const std::string &config_file) {
  26. config_map_ = LoadConfig(config_file);
  27. this->use_gpu = bool(stoi(config_map_["use_gpu"]));
  28. this->gpu_id = stoi(config_map_["gpu_id"]);
  29. this->gpu_mem = stoi(config_map_["gpu_mem"]);
  30. this->cpu_math_library_num_threads =
  31. stoi(config_map_["cpu_math_library_num_threads"]);
  32. this->use_mkldnn = bool(stoi(config_map_["use_mkldnn"]));
  33. this->max_side_len = stoi(config_map_["max_side_len"]);
  34. this->det_db_thresh = stod(config_map_["det_db_thresh"]);
  35. this->det_db_box_thresh = stod(config_map_["det_db_box_thresh"]);
  36. this->det_db_unclip_ratio = stod(config_map_["det_db_unclip_ratio"]);
  37. this->det_model_dir.assign(config_map_["det_model_dir"]);
  38. this->rec_model_dir.assign(config_map_["rec_model_dir"]);
  39. this->char_list_file.assign(config_map_["char_list_file"]);
  40. this->use_angle_cls = bool(stoi(config_map_["use_angle_cls"]));
  41. this->cls_model_dir.assign(config_map_["cls_model_dir"]);
  42. this->cls_thresh = stod(config_map_["cls_thresh"]);
  43. this->visualize = bool(stoi(config_map_["visualize"]));
  44. this->use_tensorrt = bool(stoi(config_map_["use_tensorrt"]));
  45. this->use_fp16 = bool(stod(config_map_["use_fp16"]));
  46. }
  47. bool use_gpu = false;
  48. int gpu_id = 0;
  49. int gpu_mem = 4000;
  50. int cpu_math_library_num_threads = 1;
  51. bool use_mkldnn = false;
  52. int max_side_len = 960;
  53. double det_db_thresh = 0.3;
  54. double det_db_box_thresh = 0.5;
  55. double det_db_unclip_ratio = 2.0;
  56. std::string det_model_dir;
  57. std::string rec_model_dir;
  58. bool use_angle_cls;
  59. std::string char_list_file;
  60. std::string cls_model_dir;
  61. double cls_thresh;
  62. bool visualize = true;
  63. bool use_tensorrt = false;
  64. bool use_fp16 = false;
  65. void PrintConfigInfo();
  66. private:
  67. // Load configuration
  68. std::map<std::string, std::string> LoadConfig(const std::string &config_file);
  69. std::vector<std::string> split(const std::string &str,
  70. const std::string &delim);
  71. std::map<std::string, std::string> config_map_;
  72. };
  73. } // namespace PaddleOCR