visualize.cpp 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  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/visualize.h"
  15. namespace PaddleX {
  16. std::vector<int> GenerateColorMap(int num_class) {
  17. auto colormap = std::vector<int>(3 * num_class, 0);
  18. for (int i = 0; i < num_class; ++i) {
  19. int j = 0;
  20. int lab = i;
  21. while (lab) {
  22. colormap[i * 3] |= (((lab >> 0) & 1) << (7 - j));
  23. colormap[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j));
  24. colormap[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j));
  25. ++j;
  26. lab >>= 3;
  27. }
  28. }
  29. return colormap;
  30. }
  31. cv::Mat Visualize(const cv::Mat& img,
  32. const DetResult& result,
  33. const std::map<int, std::string>& labels,
  34. float threshold) {
  35. auto colormap = GenerateColorMap(labels.size());
  36. cv::Mat vis_img = img.clone();
  37. auto boxes = result.boxes;
  38. for (int i = 0; i < boxes.size(); ++i) {
  39. if (boxes[i].score < threshold) {
  40. continue;
  41. }
  42. cv::Rect roi = cv::Rect(boxes[i].coordinate[0],
  43. boxes[i].coordinate[1],
  44. boxes[i].coordinate[2],
  45. boxes[i].coordinate[3]);
  46. // draw box and title
  47. std::string text = boxes[i].category;
  48. int c1 = colormap[3 * boxes[i].category_id + 0];
  49. int c2 = colormap[3 * boxes[i].category_id + 1];
  50. int c3 = colormap[3 * boxes[i].category_id + 2];
  51. cv::Scalar roi_color = cv::Scalar(c1, c2, c3);
  52. text += std::to_string(static_cast<int>(boxes[i].score * 100)) + "%";
  53. int font_face = cv::FONT_HERSHEY_SIMPLEX;
  54. double font_scale = 0.5f;
  55. float thickness = 0.5;
  56. cv::Size text_size =
  57. cv::getTextSize(text, font_face, font_scale, thickness, nullptr);
  58. cv::Point origin;
  59. origin.x = roi.x;
  60. origin.y = roi.y;
  61. // background
  62. cv::Rect text_back = cv::Rect(boxes[i].coordinate[0],
  63. boxes[i].coordinate[1] - text_size.height,
  64. text_size.width,
  65. text_size.height);
  66. // draw
  67. cv::rectangle(vis_img, roi, roi_color, 2);
  68. cv::rectangle(vis_img, text_back, roi_color, -1);
  69. cv::putText(vis_img,
  70. text,
  71. origin,
  72. font_face,
  73. font_scale,
  74. cv::Scalar(255, 255, 255),
  75. thickness);
  76. // mask
  77. if (boxes[i].mask.data.size() == 0) {
  78. continue;
  79. }
  80. std::vector<float> mask_data;
  81. mask_data.assign(boxes[i].mask.data.begin(), boxes[i].mask.data.end());
  82. cv::Mat bin_mask(boxes[i].mask.shape[1],
  83. boxes[i].mask.shape[0],
  84. CV_32FC1,
  85. mask_data.data());
  86. cv::Mat full_mask = cv::Mat::zeros(vis_img.size(), CV_8UC1);
  87. bin_mask.copyTo(full_mask(roi));
  88. cv::Mat mask_ch[3];
  89. mask_ch[0] = full_mask * c1;
  90. mask_ch[1] = full_mask * c2;
  91. mask_ch[2] = full_mask * c3;
  92. cv::Mat mask;
  93. cv::merge(mask_ch, 3, mask);
  94. cv::addWeighted(vis_img, 1, mask, 0.5, 0, vis_img);
  95. }
  96. return vis_img;
  97. }
  98. cv::Mat Visualize(const cv::Mat& img,
  99. const SegResult& result,
  100. const std::map<int, std::string>& labels) {
  101. auto colormap = GenerateColorMap(labels.size());
  102. std::vector<uint8_t> label_map(result.label_map.data.begin(),
  103. result.label_map.data.end());
  104. cv::Mat mask(result.label_map.shape[0],
  105. result.label_map.shape[1],
  106. CV_8UC1,
  107. label_map.data());
  108. cv::Mat color_mask = cv::Mat::zeros(
  109. result.label_map.shape[0], result.label_map.shape[1], CV_8UC3);
  110. int rows = img.rows;
  111. int cols = img.cols;
  112. for (int i = 0; i < rows; i++) {
  113. for (int j = 0; j < cols; j++) {
  114. int category_id = static_cast<int>(mask.at<uchar>(i, j));
  115. color_mask.at<cv::Vec3b>(i, j)[0] = colormap[3 * category_id + 0];
  116. color_mask.at<cv::Vec3b>(i, j)[1] = colormap[3 * category_id + 1];
  117. color_mask.at<cv::Vec3b>(i, j)[2] = colormap[3 * category_id + 2];
  118. }
  119. }
  120. return color_mask;
  121. }
  122. std::string generate_save_path(const std::string& save_dir,
  123. const std::string& file_path) {
  124. if (access(save_dir.c_str(), 0) < 0) {
  125. #ifdef _WIN32
  126. mkdir(save_dir.c_str());
  127. #else
  128. if (mkdir(save_dir.c_str(), S_IRWXU) < 0) {
  129. std::cerr << "Fail to create " << save_dir << "directory." << std::endl;
  130. }
  131. #endif
  132. }
  133. int pos = file_path.find_last_of(OS_PATH_SEP);
  134. std::string image_name(file_path.substr(pos + 1));
  135. return save_dir + OS_PATH_SEP + image_name;
  136. }
  137. } // namespace PaddleX