config_parser.h 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  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 <iostream>
  16. #include <map>
  17. #include <string>
  18. #include <vector>
  19. #include "yaml-cpp/yaml.h"
  20. #ifdef _WIN32
  21. #define OS_PATH_SEP "\\"
  22. #else
  23. #define OS_PATH_SEP "/"
  24. #endif
  25. namespace PaddleDetection {
  26. // Inference model configuration parser
  27. class ConfigPaser {
  28. public:
  29. ConfigPaser() {}
  30. ~ConfigPaser() {}
  31. bool load_config(const std::string& model_dir,
  32. const std::string& cfg = "infer_cfg.yml") {
  33. // Load as a YAML::Node
  34. YAML::Node config;
  35. config = YAML::LoadFile(model_dir + OS_PATH_SEP + cfg);
  36. // Get runtime mode : paddle, trt_fp16, trt_fp32
  37. if (config["mode"].IsDefined()) {
  38. mode_ = config["mode"].as<std::string>();
  39. } else {
  40. std::cerr << "Please set mode, "
  41. << "support value : paddle/trt_fp16/trt_fp32." << std::endl;
  42. return false;
  43. }
  44. // Get model arch : YOLO, SSD, RetinaNet, RCNN, Face
  45. if (config["arch"].IsDefined()) {
  46. arch_ = config["arch"].as<std::string>();
  47. } else {
  48. std::cerr << "Please set model arch,"
  49. << "support value : YOLO, SSD, RetinaNet, RCNN, Face."
  50. << std::endl;
  51. return false;
  52. }
  53. // Get min_subgraph_size for tensorrt
  54. if (config["min_subgraph_size"].IsDefined()) {
  55. min_subgraph_size_ = config["min_subgraph_size"].as<int>();
  56. } else {
  57. std::cerr << "Please set min_subgraph_size." << std::endl;
  58. return false;
  59. }
  60. // Get draw_threshold for visualization
  61. if (config["draw_threshold"].IsDefined()) {
  62. draw_threshold_ = config["draw_threshold"].as<float>();
  63. } else {
  64. std::cerr << "Please set draw_threshold." << std::endl;
  65. return false;
  66. }
  67. // Get Preprocess for preprocessing
  68. if (config["Preprocess"].IsDefined()) {
  69. preprocess_info_ = config["Preprocess"];
  70. } else {
  71. std::cerr << "Please set Preprocess." << std::endl;
  72. return false;
  73. }
  74. // Get label_list for visualization
  75. if (config["label_list"].IsDefined()) {
  76. label_list_ = config["label_list"].as<std::vector<std::string>>();
  77. } else {
  78. std::cerr << "Please set label_list." << std::endl;
  79. return false;
  80. }
  81. // Get use_dynamic_shape for TensorRT
  82. if (config["use_dynamic_shape"].IsDefined()) {
  83. use_dynamic_shape_ = config["use_dynamic_shape"].as<bool>();
  84. } else {
  85. std::cerr << "Please set use_dynamic_shape." << std::endl;
  86. return false;
  87. }
  88. // Get conf_thresh for tracker
  89. if (config["tracker"].IsDefined()) {
  90. if (config["tracker"]["conf_thres"].IsDefined()) {
  91. conf_thresh_ = config["tracker"]["conf_thres"].as<float>();
  92. } else {
  93. std::cerr << "Please set conf_thres in tracker." << std::endl;
  94. return false;
  95. }
  96. }
  97. // Get NMS for postprocess
  98. if (config["NMS"].IsDefined()) {
  99. nms_info_ = config["NMS"];
  100. }
  101. // Get fpn_stride in PicoDet
  102. if (config["fpn_stride"].IsDefined()) {
  103. fpn_stride_.clear();
  104. for (auto item : config["fpn_stride"]) {
  105. fpn_stride_.emplace_back(item.as<int>());
  106. }
  107. }
  108. if (config["mask"].IsDefined()) {
  109. mask_ = config["mask"].as<bool>();
  110. }
  111. return true;
  112. }
  113. std::string mode_;
  114. float draw_threshold_;
  115. std::string arch_;
  116. int min_subgraph_size_;
  117. YAML::Node preprocess_info_;
  118. YAML::Node nms_info_;
  119. std::vector<std::string> label_list_;
  120. std::vector<int> fpn_stride_;
  121. bool use_dynamic_shape_;
  122. float conf_thresh_;
  123. bool mask_ = false;
  124. };
  125. } // namespace PaddleDetection