cnosd.cpp 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216
  1. /*************************************************************************
  2. * Copyright (C) [2019] by Cambricon, Inc. All rights reserved
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * The above copyright notice and this permission notice shall be included in
  11. * all copies or substantial portions of the Software.
  12. * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
  13. * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  14. * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
  15. * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  16. * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  17. * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  18. * THE SOFTWARE.
  19. *************************************************************************/
  20. #include "cnosd.h"
  21. #include <algorithm>
  22. #include <string>
  23. #include <vector>
  24. using std::to_string;
  25. #if CV_VERSION_EPOCH == 2
  26. #define OPENCV_MAJOR_VERSION 2
  27. #elif CV_VERSION_MAJOR >= 3
  28. #define OPENCV_MAJOR_VERSION CV_VERSION_MAJOR
  29. #endif
  30. // Keep 2 digits after decimal
  31. static string FloatToString(float number) {
  32. char buffer[10];
  33. snprintf(buffer, sizeof(buffer), "%.2f", number);
  34. return string(buffer);
  35. }
  36. // http://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically
  37. static Scalar HSV2RGB(const float h, const float s, const float v) {
  38. const int h_i = static_cast<int>(h * 6);
  39. const float f = h * 6 - h_i;
  40. const float p = v * (1 - s);
  41. const float q = v * (1 - f * s);
  42. const float t = v * (1 - (1 - f) * s);
  43. float r, g, b;
  44. switch (h_i) {
  45. case 0:
  46. r = v;
  47. g = t;
  48. b = p;
  49. break;
  50. case 1:
  51. r = q;
  52. g = v;
  53. b = p;
  54. break;
  55. case 2:
  56. r = p;
  57. g = v;
  58. b = t;
  59. break;
  60. case 3:
  61. r = p;
  62. g = q;
  63. b = v;
  64. break;
  65. case 4:
  66. r = t;
  67. g = p;
  68. b = v;
  69. break;
  70. case 5:
  71. r = v;
  72. g = p;
  73. b = q;
  74. break;
  75. default:
  76. r = 1;
  77. g = 1;
  78. b = 1;
  79. break;
  80. }
  81. return Scalar(r * 255, g * 255, b * 255);
  82. }
  83. static vector<Scalar> GenerateColors(const int n) {
  84. vector<Scalar> colors;
  85. cv::RNG rng(12345);
  86. const float golden_ratio_conjugate = 0.618033988749895f;
  87. const float s = 0.3f;
  88. const float v = 0.99f;
  89. for (int i = 0; i < n; ++i) {
  90. const float h = std::fmod(rng.uniform(0.0f, 1.0f) + golden_ratio_conjugate, 1.0f);
  91. colors.push_back(HSV2RGB(h, s, v));
  92. }
  93. return colors;
  94. }
  95. static vector<string> LoadLabels(const string& filename) {
  96. vector<string> labels;
  97. std::ifstream file(filename);
  98. if (file.is_open()) {
  99. string line;
  100. while (std::getline(file, line)) {
  101. labels.push_back(string(line));
  102. }
  103. file.close();
  104. } else {
  105. printf("[Warning]: Load labels failed: %s\n", filename.c_str());
  106. }
  107. return labels;
  108. }
  109. CnOsd::CnOsd(const vector<string>& labels) : labels_(labels) {
  110. colors_ = ::GenerateColors(labels_.size());
  111. }
  112. CnOsd::CnOsd(const string& label_fname) {
  113. LoadLabels(label_fname);
  114. }
  115. void CnOsd::LoadLabels(const std::string& fname) {
  116. labels_ = ::LoadLabels(fname);
  117. colors_ = ::GenerateColors(labels_.size());
  118. }
  119. void CnOsd::set_font(int font) { font_ = font; }
  120. void CnOsd::DrawId(Mat image, string text) const {
  121. float scale = CalScale(image.cols * image.rows);
  122. cv::Size text_size = cv::getTextSize(text, font_, scale, 1, nullptr);
  123. Scalar color(0, 255, 255);
  124. cv::putText(image, text, Point(0, text_size.height), font_, scale, color, 1, 8, false);
  125. }
  126. void CnOsd::DrawFps(Mat image, float fps) const {
  127. // check input data
  128. if (image.cols * image.rows == 0) {
  129. return;
  130. }
  131. string text = "fps: " + ::FloatToString(fps);
  132. float scale = CalScale(image.cols * image.rows);
  133. cv::Size text_size = cv::getTextSize(text, font_, scale, 1, nullptr);
  134. Scalar color(0, 0, 255);
  135. cv::putText(image, text, Point(image.cols - text_size.width, text_size.height), font_, scale, color, 1, 8, false);
  136. }
  137. // tl: top left
  138. // br: bottom right
  139. // bl: bottom left
  140. void CnOsd::DrawLabel(Mat image, const vector<edk::DetectObject>& objects) const {
  141. // check input data
  142. if (image.rows * image.cols == 0) {
  143. return;
  144. }
  145. for (auto& object : objects) {
  146. float xmin = object.bbox.x * image.cols;
  147. float ymin = object.bbox.y * image.rows;
  148. float xmax = (object.bbox.x + object.bbox.width) * image.cols;
  149. float ymax = (object.bbox.y + object.bbox.height) * image.rows;
  150. string text;
  151. Scalar color;
  152. if (labels().size() <= static_cast<size_t>(object.label)) {
  153. text = "Label not found, id = " + to_string(object.label);
  154. color = Scalar(0, 0, 0);
  155. } else {
  156. text = labels()[object.label];
  157. color = colors_[object.label];
  158. }
  159. // Detection window
  160. Point tl(xmin, ymin);
  161. Point br(xmax, ymax);
  162. int box_thickness = get_box_thickness();
  163. cv::rectangle(image, tl, br, color, box_thickness);
  164. // Label and Score
  165. text += " " + FloatToString(object.score);
  166. // Track Id
  167. if (object.track_id >= 0) text += " track_id:" + to_string(object.track_id);
  168. float scale = CalScale(image.cols * image.rows);
  169. int text_thickness = 1;
  170. cv::Size text_size = cv::getTextSize(text, font_, scale, text_thickness, nullptr);
  171. int offset = (box_thickness == 1 ? 0 : -(box_thickness + 1) / 2);
  172. Point bl(xmin + offset, ymax + offset);
  173. Point label_left, label_right;
  174. label_left = bl;
  175. label_right = bl + Point(text_size.width + offset, text_size.height * 1.4);
  176. if (label_right.y > image.rows) {
  177. label_right.y -= text_size.height * 1.4;
  178. label_left.y -= text_size.height * 1.4;
  179. }
  180. if (label_right.x > image.cols) {
  181. label_right.x = image.cols;
  182. label_left.x = image.cols - text_size.width;
  183. }
  184. #if OPENCV_MAJOR_VERSION > 2
  185. cv::rectangle(image, label_left, label_right, color, cv::FILLED);
  186. #else
  187. cv::rectangle(image, label_left, label_right, color, CV_FILLED);
  188. #endif
  189. cv::putText(image, text, label_left + Point(0, text_size.height), font_, scale, Scalar(255, 255, 255) - color,
  190. text_thickness, 8, false);
  191. }
  192. }