paddlex.cpp 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256
  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. #include <iostream>
  16. #include <fstream>
  17. namespace PaddleX {
  18. void Model::create_predictor(const std::string& model_dir,
  19. const std::string& cfg_file,
  20. int thread_num) {
  21. paddle::lite_api::MobileConfig config;
  22. config.set_model_from_file(model_dir);
  23. config.set_threads(thread_num);
  24. load_config(cfg_file);
  25. predictor_ =
  26. paddle::lite_api::CreatePaddlePredictor<paddle::lite_api::MobileConfig>(
  27. config);
  28. }
  29. bool Model::load_config(const std::string& cfg_file) {
  30. YAML::Node config = YAML::LoadFile(cfg_file);
  31. type = config["_Attributes"]["model_type"].as<std::string>();
  32. name = config["Model"].as<std::string>();
  33. bool to_rgb = true;
  34. if (config["TransformsMode"].IsDefined()) {
  35. std::string mode = config["TransformsMode"].as<std::string>();
  36. if (mode == "BGR") {
  37. to_rgb = false;
  38. } else if (mode != "RGB") {
  39. std::cerr << "[Init] Only 'RGB' or 'BGR' is supported for TransformsMode"
  40. << std::endl;
  41. return false;
  42. }
  43. }
  44. // init preprocess ops
  45. transforms_.Init(config["Transforms"], to_rgb);
  46. // read label list
  47. for (const auto& item : config["_Attributes"]["labels"]) {
  48. int index = labels.size();
  49. labels[index] = item.as<std::string>();
  50. }
  51. return true;
  52. }
  53. bool Model::preprocess(cv::Mat* input_im, ImageBlob* inputs) {
  54. if (!transforms_.Run(input_im, inputs)) {
  55. return false;
  56. }
  57. return true;
  58. }
  59. bool Model::predict(const cv::Mat& im, ClsResult* result) {
  60. inputs_.clear();
  61. if (type == "detector") {
  62. std::cerr << "Loading model is a 'detector', DetResult should be passed to "
  63. "function predict()!"
  64. << std::endl;
  65. return false;
  66. } else if (type == "segmenter") {
  67. std::cerr << "Loading model is a 'segmenter', SegResult should be passed "
  68. "to function predict()!"
  69. << std::endl;
  70. return false;
  71. }
  72. // preprocess
  73. inputs_.input_tensor_ = std::move(predictor_->GetInput(0));
  74. cv::Mat im_clone = im.clone();
  75. if (!preprocess(&im_clone, &inputs_)) {
  76. std::cerr << "Preprocess failed!" << std::endl;
  77. return false;
  78. }
  79. // predict
  80. predictor_->Run();
  81. std::unique_ptr<const paddle::lite_api::Tensor> output_tensor(
  82. std::move(predictor_->GetOutput(0)));
  83. const float *outputs_data = output_tensor->mutable_data<float>();
  84. auto output_shape = output_tensor->shape();
  85. int64_t size = 1;
  86. for (const auto& i : output_shape) {
  87. size *= i;
  88. }
  89. // postprocess
  90. auto ptr = std::max_element(outputs_data, outputs_data + size);
  91. result->category_id = std::distance(outputs_data, ptr);
  92. result->score = *ptr;
  93. result->category = labels[result->category_id];
  94. }
  95. bool Model::predict(const cv::Mat& im, DetResult* result) {
  96. inputs_.clear();
  97. result->clear();
  98. if (type == "classifier") {
  99. std::cerr << "Loading model is a 'classifier', ClsResult should be passed "
  100. "to function predict()!" << std::endl;
  101. return false;
  102. } else if (type == "segmenter") {
  103. std::cerr << "Loading model is a 'segmenter', SegResult should be passed "
  104. "to function predict()!" << std::endl;
  105. return false;
  106. }
  107. inputs_.input_tensor_ = std::move(predictor_->GetInput(0));
  108. cv::Mat im_clone = im.clone();
  109. if (!preprocess(&im_clone, &inputs_)) {
  110. std::cerr << "Preprocess failed!" << std::endl;
  111. return false;
  112. }
  113. if (name == "YOLOv3") {
  114. std::unique_ptr<paddle::lite_api::Tensor> im_size_tensor(
  115. std::move(predictor_->GetInput(1)));
  116. im_size_tensor->Resize({1, 2});
  117. auto *im_size_data = im_size_tensor->mutable_data<int>();
  118. memcpy(im_size_data, inputs_.ori_im_size_.data(), 1*2*sizeof(int));
  119. }
  120. predictor_->Run();
  121. auto output_names = predictor_->GetOutputNames();
  122. auto output_box_tensor = predictor_->GetTensor(output_names[0]);
  123. auto *output_box = output_box_tensor->mutable_data<float>();
  124. auto output_box_shape = output_box_tensor->shape();
  125. int64_t size = 1;
  126. for (const auto& i : output_box_shape) {
  127. size *= i;
  128. }
  129. auto num_boxes = static_cast<int>(size / 6);
  130. for (int i = 0; i < num_boxes; ++i) {
  131. Box box;
  132. box.category_id = static_cast<int>(round(output_box[i * 6]));
  133. box.category = labels[box.category_id];
  134. box.score = output_box[i * 6 + 1];
  135. float xmin = output_box[i * 6 + 2];
  136. float ymin = output_box[i * 6 + 3];
  137. float xmax = output_box[i * 6 + 4];
  138. float ymax = output_box[i * 6 + 5];
  139. float w = xmax - xmin + 1;
  140. float h = ymax - ymin + 1;
  141. box.coordinate = {xmin, ymin, w, h};
  142. result->boxes.push_back(std::move(box));
  143. }
  144. return true;
  145. }
  146. bool Model::predict(const cv::Mat& im, SegResult* result) {
  147. result->clear();
  148. inputs_.clear();
  149. if (type == "classifier") {
  150. std::cerr << "Loading model is a 'classifier', ClsResult should be passed "
  151. "to function predict()!" << std::endl;
  152. return false;
  153. } else if (type == "detector") {
  154. std::cerr << "Loading model is a 'detector', DetResult should be passed to "
  155. "function predict()!" << std::endl;
  156. return false;
  157. }
  158. inputs_.input_tensor_ = std::move(predictor_->GetInput(0));
  159. cv::Mat im_clone = im.clone();
  160. if (!preprocess(&im_clone, &inputs_)) {
  161. std::cerr << "Preprocess failed!" << std::endl;
  162. return false;
  163. }
  164. std::cout << "Preprocess is done" << std::endl;
  165. predictor_->Run();
  166. auto output_names = predictor_->GetOutputNames();
  167. auto output_label_tensor = predictor_->GetTensor(output_names[0]);
  168. const int64_t *label_data = output_label_tensor->mutable_data<int64_t>();
  169. std::vector<int64_t> output_label_shape = output_label_tensor->shape();
  170. int size = 1;
  171. for (const auto& i : output_label_shape) {
  172. size *= i;
  173. result->label_map.shape.push_back(i);
  174. }
  175. result->label_map.data.resize(size);
  176. memcpy(result->label_map.data.data(), label_data, size*sizeof(int64_t));
  177. auto output_score_tensor = predictor_->GetTensor(output_names[1]);
  178. const float *score_data = output_score_tensor->mutable_data<float>();
  179. std::vector<int64_t> output_score_shape = output_score_tensor->shape();
  180. size = 1;
  181. for (const auto& i : output_score_shape) {
  182. size *= i;
  183. result->score_map.shape.push_back(i);
  184. }
  185. result->score_map.data.resize(size);
  186. memcpy(result->score_map.data.data(), score_data, size*sizeof(float));
  187. std::vector<uint8_t> label_map(result->label_map.data.begin(),
  188. result->label_map.data.end());
  189. cv::Mat mask_label(result->label_map.shape[1],
  190. result->label_map.shape[2],
  191. CV_8UC1,
  192. label_map.data());
  193. cv::Mat mask_score(result->score_map.shape[2],
  194. result->score_map.shape[3],
  195. CV_32FC1,
  196. result->score_map.data.data());
  197. int idx = 1;
  198. int len_postprocess = inputs_.im_size_before_resize_.size();
  199. for (std::vector<std::string>::reverse_iterator iter =
  200. inputs_.reshape_order_.rbegin();
  201. iter != inputs_.reshape_order_.rend();
  202. ++iter) {
  203. if (*iter == "padding") {
  204. auto before_shape = inputs_.im_size_before_resize_[len_postprocess - idx];
  205. inputs_.im_size_before_resize_.pop_back();
  206. auto padding_w = before_shape[0];
  207. auto padding_h = before_shape[1];
  208. mask_label = mask_label(cv::Rect(0, 0, padding_h, padding_w));
  209. mask_score = mask_score(cv::Rect(0, 0, padding_h, padding_w));
  210. } else if (*iter == "resize") {
  211. auto before_shape = inputs_.im_size_before_resize_[len_postprocess - idx];
  212. inputs_.im_size_before_resize_.pop_back();
  213. auto resize_w = before_shape[0];
  214. auto resize_h = before_shape[1];
  215. cv::resize(mask_label,
  216. mask_label,
  217. cv::Size(resize_h, resize_w),
  218. 0,
  219. 0,
  220. cv::INTER_NEAREST);
  221. cv::resize(mask_score,
  222. mask_score,
  223. cv::Size(resize_h, resize_w),
  224. 0,
  225. 0,
  226. cv::INTER_LINEAR);
  227. }
  228. ++idx;
  229. }
  230. result->label_map.data.assign(mask_label.begin<uint8_t>(),
  231. mask_label.end<uint8_t>());
  232. result->label_map.shape = {mask_label.rows, mask_label.cols};
  233. result->score_map.data.assign(mask_score.begin<float>(),
  234. mask_score.end<float>());
  235. result->score_map.shape = {mask_score.rows, mask_score.cols};
  236. return true;
  237. }
  238. } // namespace PaddleX