seg_postprocess.cpp 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. // Copyright (c) 2021 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 "model_deploy/ppseg/include/seg_postprocess.h"
  15. #include <time.h>
  16. namespace PaddleDeploy {
  17. bool SegPostprocess::Init(const YAML::Node& yaml_config) {
  18. if (yaml_config["version"].IsDefined() &&
  19. yaml_config["toolkit"].as<std::string>() == "PaddleX") {
  20. version_ = yaml_config["version"].as<std::string>();
  21. } else {
  22. version_ = "0.0.0";
  23. }
  24. return true;
  25. }
  26. void SegPostprocess::RestoreSegMap(const ShapeInfo& shape_info,
  27. cv::Mat* label_mat, cv::Mat* score_mat,
  28. SegResult* result) {
  29. int ori_h = shape_info.shapes[0][1];
  30. int ori_w = shape_info.shapes[0][0];
  31. int score_c = score_mat->channels();
  32. result->label_map.Resize({ori_h, ori_w});
  33. result->score_map.Resize({ori_h, ori_w, score_c});
  34. for (int j = shape_info.transforms.size() - 1; j > 0; --j) {
  35. std::vector<int> last_shape = shape_info.shapes[j - 1];
  36. std::vector<int> cur_shape = shape_info.shapes[j];
  37. if (shape_info.transforms[j] == "Resize" ||
  38. shape_info.transforms[j] == "ResizeByShort" ||
  39. shape_info.transforms[j] == "ResizeByLong") {
  40. if (last_shape[0] != label_mat->cols ||
  41. last_shape[1] != label_mat->rows) {
  42. cv::resize(*label_mat, *label_mat,
  43. cv::Size(last_shape[0], last_shape[1]), 0, 0,
  44. cv::INTER_NEAREST);
  45. cv::resize(*score_mat, *score_mat,
  46. cv::Size(last_shape[0], last_shape[1]), 0, 0,
  47. cv::INTER_LINEAR);
  48. }
  49. } else if (shape_info.transforms[j] == "Padding") {
  50. if (last_shape[0] < label_mat->cols || last_shape[1] < label_mat->rows) {
  51. *label_mat = (*label_mat)(cv::Rect(0, 0, last_shape[0], last_shape[1]));
  52. *score_mat = (*score_mat)(cv::Rect(0, 0, last_shape[0], last_shape[1]));
  53. }
  54. }
  55. }
  56. if (label_mat->isContinuous()) {
  57. result->label_map.data.assign(
  58. reinterpret_cast<const uint8_t*>(label_mat->data),
  59. reinterpret_cast<const uint8_t*>(label_mat->data) +
  60. label_mat->total() * (label_mat->channels()));
  61. } else {
  62. for (int i = 0; i < label_mat->rows; ++i) {
  63. result->label_map.data.insert(
  64. result->label_map.data.end(), label_mat->ptr<uint8_t>(i),
  65. label_mat->ptr<uint8_t>(i) +
  66. label_mat->cols * (label_mat->channels()));
  67. }
  68. }
  69. if (score_mat->isContinuous()) {
  70. result->score_map.data.assign(
  71. reinterpret_cast<const float*>(score_mat->data),
  72. reinterpret_cast<const float*>(score_mat->data) +
  73. score_mat->total() * (score_mat->channels()));
  74. } else {
  75. for (int i = 0; i < score_mat->rows; ++i) {
  76. result->score_map.data.insert(
  77. result->score_map.data.end(), score_mat->ptr<float>(i),
  78. score_mat->ptr<float>(i) + score_mat->cols * (score_mat->channels()));
  79. }
  80. }
  81. }
  82. // ppseg version >= 2.1 shape = [b, w, h]
  83. bool SegPostprocess::RunV2(const DataBlob& output,
  84. const std::vector<ShapeInfo>& shape_infos,
  85. std::vector<Result>* results, int thread_num) {
  86. int batch_size = shape_infos.size();
  87. int label_map_size = output.shape[1] * output.shape[2];
  88. const uint8_t* label_data;
  89. std::vector<uint8_t> label_vector;
  90. if (output.dtype == INT64) { // int64
  91. const int64_t* output_data =
  92. reinterpret_cast<const int64_t*>(output.data.data());
  93. std::transform(output_data, output_data + label_map_size * batch_size,
  94. std::back_inserter(label_vector),
  95. [](int64_t x) { return (uint8_t)x; });
  96. label_data = reinterpret_cast<const uint8_t*>(label_vector.data());
  97. } else if (output.dtype == INT32) { // int32
  98. const int32_t* output_data =
  99. reinterpret_cast<const int32_t*>(output.data.data());
  100. std::transform(output_data, output_data + label_map_size * batch_size,
  101. std::back_inserter(label_vector),
  102. [](int32_t x) { return (uint8_t)x; });
  103. label_data = reinterpret_cast<const uint8_t*>(label_vector.data());
  104. } else if (output.dtype == INT8) { // uint8
  105. label_data = reinterpret_cast<const uint8_t*>(output.data.data());
  106. } else {
  107. std::cerr << "Output dtype is not support on seg posrtprocess "
  108. << output.dtype << std::endl;
  109. return false;
  110. }
  111. for (int i = 0; i < batch_size; ++i) {
  112. (*results)[i].model_type = "seg";
  113. (*results)[i].seg_result = new SegResult();
  114. const uint8_t* current_start_ptr = label_data + i * label_map_size;
  115. cv::Mat score_mat(output.shape[1], output.shape[2], CV_32FC1,
  116. cv::Scalar(1.0));
  117. cv::Mat label_mat(output.shape[1], output.shape[2], CV_8UC1,
  118. const_cast<uint8_t*>(current_start_ptr));
  119. RestoreSegMap(shape_infos[i], &label_mat, &score_mat,
  120. (*results)[i].seg_result);
  121. }
  122. return true;
  123. }
  124. // paddlex version >= 2.0.0 shape = [b, h, w, c]
  125. bool SegPostprocess::RunXV2(const std::vector<DataBlob>& outputs,
  126. const std::vector<ShapeInfo>& shape_infos,
  127. std::vector<Result>* results, int thread_num) {
  128. int batch_size = shape_infos.size();
  129. int label_map_size = outputs[0].shape[1] * outputs[0].shape[2];
  130. std::vector<int> score_map_shape = outputs[1].shape;
  131. int score_map_size =
  132. std::accumulate(score_map_shape.begin() + 1, score_map_shape.end(), 1,
  133. std::multiplies<int>());
  134. const uint8_t* label_map_data;
  135. std::vector<uint8_t> label_map_vector;
  136. if (outputs[0].dtype == INT32) {
  137. const int32_t* output_data =
  138. reinterpret_cast<const int32_t*>(outputs[0].data.data());
  139. std::transform(output_data, output_data + label_map_size * batch_size,
  140. std::back_inserter(label_map_vector),
  141. [](int32_t x) { return (uint8_t)x; });
  142. label_map_data = reinterpret_cast<const uint8_t*>(label_map_vector.data());
  143. }
  144. const float* score_map_data =
  145. reinterpret_cast<const float*>(outputs[1].data.data());
  146. for (int i = 0; i < batch_size; ++i) {
  147. (*results)[i].model_type = "seg";
  148. (*results)[i].seg_result = new SegResult();
  149. const uint8_t* current_label_start_ptr =
  150. label_map_data + i * label_map_size;
  151. const float* current_score_start_ptr = score_map_data + i * score_map_size;
  152. cv::Mat label_mat(outputs[0].shape[1], outputs[0].shape[2], CV_8UC1,
  153. const_cast<uint8_t*>(current_label_start_ptr));
  154. cv::Mat score_mat(score_map_shape[1], score_map_shape[2],
  155. CV_32FC(score_map_shape[3]),
  156. const_cast<float*>(current_score_start_ptr));
  157. RestoreSegMap(shape_infos[i], &label_mat, &score_mat,
  158. (*results)[i].seg_result);
  159. }
  160. return true;
  161. }
  162. bool SegPostprocess::Run(const std::vector<DataBlob>& outputs,
  163. const std::vector<ShapeInfo>& shape_infos,
  164. std::vector<Result>* results, int thread_num) {
  165. if (outputs.size() == 0) {
  166. std::cerr << "empty output on SegPostprocess" << std::endl;
  167. return true;
  168. }
  169. results->clear();
  170. int batch_size = shape_infos.size();
  171. results->resize(batch_size);
  172. // tricks for PaddleX, of which segmentation model has two outputs
  173. int index = 0;
  174. if (outputs.size() == 2) {
  175. index = 1;
  176. }
  177. std::vector<int> score_map_shape = outputs[index].shape;
  178. // paddlex version >= 2.0.0 shape[b, h, w, c]
  179. if (version_ >= "2.0.0") {
  180. return RunXV2(outputs, shape_infos, results, thread_num);
  181. }
  182. // ppseg version >= 2.1 shape = [b, h, w]
  183. if (score_map_shape.size() == 3) {
  184. return RunV2(outputs[index], shape_infos, results, thread_num);
  185. }
  186. int score_map_size =
  187. std::accumulate(score_map_shape.begin() + 1, score_map_shape.end(), 1,
  188. std::multiplies<int>());
  189. const float* score_map_data =
  190. reinterpret_cast<const float*>(outputs[index].data.data());
  191. int num_map_pixels = score_map_shape[2] * score_map_shape[3];
  192. for (int i = 0; i < batch_size; ++i) {
  193. (*results)[i].model_type = "seg";
  194. (*results)[i].seg_result = new SegResult();
  195. const float* current_start_ptr = score_map_data + i * score_map_size;
  196. cv::Mat ori_score_mat(score_map_shape[1],
  197. score_map_shape[2] * score_map_shape[3], CV_32FC1,
  198. const_cast<float*>(current_start_ptr));
  199. ori_score_mat = ori_score_mat.t();
  200. cv::Mat score_mat(score_map_shape[2], score_map_shape[3], CV_32FC1);
  201. cv::Mat label_mat(score_map_shape[2], score_map_shape[3], CV_8UC1);
  202. for (int j = 0; j < ori_score_mat.rows; ++j) {
  203. double max_value;
  204. cv::Point max_id;
  205. minMaxLoc(ori_score_mat.row(j), 0, &max_value, 0, &max_id);
  206. score_mat.at<float>(j) = max_value;
  207. label_mat.at<uchar>(j) = max_id.x;
  208. }
  209. RestoreSegMap(shape_infos[i], &label_mat, &score_mat,
  210. (*results)[i].seg_result);
  211. }
  212. return true;
  213. }
  214. } // namespace PaddleDeploy