paddlex.cpp 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  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 "include/paddlex/paddlex.h"
  15. using namespace InferenceEngine;
  16. namespace PaddleX {
  17. void Model::create_predictor(const std::string& model_dir,
  18. const std::string& cfg_dir,
  19. std::string device) {
  20. Core ie;
  21. network_ = ie.ReadNetwork(model_dir, model_dir.substr(0, model_dir.size() - 4) + ".bin");
  22. network_.setBatchSize(1);
  23. InputInfo::Ptr input_info = network_.getInputsInfo().begin()->second;
  24. input_info->getPreProcess().setResizeAlgorithm(RESIZE_BILINEAR);
  25. input_info->setLayout(Layout::NCHW);
  26. input_info->setPrecision(Precision::FP32);
  27. executable_network_ = ie.LoadNetwork(network_, device);
  28. load_config(cfg_dir);
  29. }
  30. bool Model::load_config(const std::string& cfg_dir) {
  31. YAML::Node config = YAML::LoadFile(cfg_dir);
  32. type = config["_Attributes"]["model_type"].as<std::string>();
  33. name = config["Model"].as<std::string>();
  34. bool to_rgb = true;
  35. if (config["TransformsMode"].IsDefined()) {
  36. std::string mode = config["TransformsMode"].as<std::string>();
  37. if (mode == "BGR") {
  38. to_rgb = false;
  39. } else if (mode != "RGB") {
  40. std::cerr << "[Init] Only 'RGB' or 'BGR' is supported for TransformsMode"
  41. << std::endl;
  42. return false;
  43. }
  44. }
  45. // 构建数据处理流
  46. transforms_.Init(config["Transforms"], to_rgb);
  47. // 读入label list
  48. labels.clear();
  49. labels = config["_Attributes"]["labels"].as<std::vector<std::string>>();
  50. return true;
  51. }
  52. bool Model::preprocess(cv::Mat* input_im) {
  53. if (!transforms_.Run(input_im, inputs_)) {
  54. return false;
  55. }
  56. return true;
  57. }
  58. bool Model::predict(const cv::Mat& im, ClsResult* result) {
  59. if (type == "detector") {
  60. std::cerr << "Loading model is a 'detector', DetResult should be passed to "
  61. "function predict()!"
  62. << std::endl;
  63. return false;
  64. } else if (type == "segmenter") {
  65. std::cerr << "Loading model is a 'segmenter', SegResult should be passed "
  66. "to function predict()!"
  67. << std::endl;
  68. return false;
  69. }
  70. // 处理输入图像
  71. InferRequest infer_request = executable_network_.CreateInferRequest();
  72. std::string input_name = network_.getInputsInfo().begin()->first;
  73. inputs_ = infer_request.GetBlob(input_name);
  74. auto im_clone = im.clone();
  75. if (!preprocess(&im_clone)) {
  76. std::cerr << "Preprocess failed!" << std::endl;
  77. return false;
  78. }
  79. infer_request.Infer();
  80. std::string output_name = network_.getOutputsInfo().begin()->first;
  81. output_ = infer_request.GetBlob(output_name);
  82. MemoryBlob::CPtr moutput = as<MemoryBlob>(output_);
  83. auto moutputHolder = moutput->rmap();
  84. float* outputs_data = moutputHolder.as<float *>();
  85. // 对模型输出结果进行后处理
  86. auto ptr = std::max_element(outputs_data, outputs_data+sizeof(outputs_data));
  87. result->category_id = std::distance(outputs_data, ptr);
  88. result->score = *ptr;
  89. result->category = labels[result->category_id];
  90. //for (int i=0;i<sizeof(outputs_data);i++){
  91. // std::cout << labels[i] << std::endl;
  92. // std::cout << outputs_[i] << std::endl;
  93. // }
  94. }
  95. } // namespce of PaddleX