yolov8_utils.cpp 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278
  1. #pragma once
  2. #include "yolov8_utils.h"
  3. //using namespace cv;
  4. //using namespace std;
  5. bool CheckParams(int netHeight, int netWidth, const int* netStride, int strideSize) {
  6. if (netHeight % netStride[strideSize - 1] != 0 || netWidth % netStride[strideSize - 1] != 0)
  7. {
  8. std::cout << "Error:_netHeight and _netWidth must be multiple of max stride " << netStride[strideSize - 1] << "!" << std::endl;
  9. return false;
  10. }
  11. return true;
  12. }
  13. bool CheckModelPath(std::string modelPath) {
  14. #ifdef _WIN32
  15. if (0 != _access(modelPath.c_str(), 0)) {
  16. std::cout << "Model path does not exist, please check " << modelPath << std::endl;
  17. return false;
  18. }
  19. else
  20. return true;
  21. #else
  22. return (0 == access(modelPath.c_str(), F_OK));
  23. if (0 != access(modelPath.c_str(), F_OK)) {
  24. std::cout << "Model path does not exist, please check " << modelPath << std::endl;
  25. return false;
  26. }
  27. else
  28. return true;
  29. #endif
  30. }
  31. void LetterBox(const cv::Mat& image, cv::Mat& outImage, cv::Vec4d& params, const cv::Size& newShape,
  32. bool autoShape, bool scaleFill, bool scaleUp, int stride, const cv::Scalar& color)
  33. {
  34. if (false) {
  35. int maxLen = MAX(image.rows, image.cols);
  36. outImage = cv::Mat::zeros(cv::Size(maxLen, maxLen), CV_8UC3);
  37. image.copyTo(outImage(cv::Rect(0, 0, image.cols, image.rows)));
  38. params[0] = 1;
  39. params[1] = 1;
  40. params[3] = 0;
  41. params[2] = 0;
  42. }
  43. cv::Size shape = image.size();
  44. float r = std::min((float)newShape.height / (float)shape.height,
  45. (float)newShape.width / (float)shape.width);
  46. if (!scaleUp)
  47. r = std::min(r, 1.0f);
  48. float ratio[2]{ r, r };
  49. int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };
  50. auto dw = (float)(newShape.width - new_un_pad[0]);
  51. auto dh = (float)(newShape.height - new_un_pad[1]);
  52. if (autoShape)
  53. {
  54. dw = (float)((int)dw % stride);
  55. dh = (float)((int)dh % stride);
  56. }
  57. else if (scaleFill)
  58. {
  59. dw = 0.0f;
  60. dh = 0.0f;
  61. new_un_pad[0] = newShape.width;
  62. new_un_pad[1] = newShape.height;
  63. ratio[0] = (float)newShape.width / (float)shape.width;
  64. ratio[1] = (float)newShape.height / (float)shape.height;
  65. }
  66. dw /= 2.0f;
  67. dh /= 2.0f;
  68. if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
  69. {
  70. cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
  71. }
  72. else {
  73. outImage = image.clone();
  74. }
  75. int top = int(std::round(dh - 0.1f));
  76. int bottom = int(std::round(dh + 0.1f));
  77. int left = int(std::round(dw - 0.1f));
  78. int right = int(std::round(dw + 0.1f));
  79. params[0] = ratio[0];
  80. params[1] = ratio[1];
  81. params[2] = left;
  82. params[3] = top;
  83. cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
  84. }
  85. void GetMask(const cv::Mat& maskProposals, const cv::Mat& maskProtos, std::vector<OutputParams>& output, const MaskParams& maskParams) {
  86. //std::cout << maskProtos.size << std::endl;
  87. int net_width = maskParams.netWidth;
  88. int net_height = maskParams.netHeight;
  89. int seg_channels = maskProtos.size[1];
  90. int seg_height = maskProtos.size[2];
  91. int seg_width = maskProtos.size[3];
  92. float mask_threshold = maskParams.maskThreshold;
  93. cv::Vec4f params = maskParams.params;
  94. cv::Size src_img_shape = maskParams.srcImgShape;
  95. cv::Mat protos = maskProtos.reshape(0, { seg_channels,seg_width * seg_height });
  96. cv::Mat matmul_res = (maskProposals * protos).t();
  97. cv::Mat masks = matmul_res.reshape(output.size(), { seg_width,seg_height });
  98. std::vector<cv::Mat> maskChannels;
  99. split(masks, maskChannels);
  100. for (int i = 0; i < output.size(); ++i) {
  101. cv::Mat dest, mask;
  102. //sigmoid
  103. cv::exp(-maskChannels[i], dest);
  104. dest = 1.0 / (1.0 + dest);
  105. cv::Rect roi(int(params[2] / net_width * seg_width), int(params[3] / net_height * seg_height), int(seg_width - params[2] / 2), int(seg_height - params[3] / 2));
  106. dest = dest(roi);
  107. resize(dest, mask, src_img_shape, cv::INTER_NEAREST);
  108. //crop
  109. cv::Rect temp_rect = output[i].box;
  110. mask = mask(temp_rect) > mask_threshold;
  111. output[i].boxMask = mask;
  112. }
  113. }
  114. void GetMask2(const cv::Mat& maskProposals, const cv::Mat& maskProtos, OutputParams& output, const MaskParams& maskParams) {
  115. int net_width = maskParams.netWidth;
  116. int net_height = maskParams.netHeight;
  117. int seg_channels = maskProtos.size[1];
  118. int seg_height = maskProtos.size[2];
  119. int seg_width = maskProtos.size[3];
  120. float mask_threshold = maskParams.maskThreshold;
  121. cv::Vec4f params = maskParams.params;
  122. cv::Size src_img_shape = maskParams.srcImgShape;
  123. cv::Rect temp_rect = output.box;
  124. //crop from mask_protos
  125. int rang_x = floor((temp_rect.x * params[0] + params[2]) / net_width * seg_width);
  126. int rang_y = floor((temp_rect.y * params[1] + params[3]) / net_height * seg_height);
  127. int rang_w = ceil(((temp_rect.x + temp_rect.width) * params[0] + params[2]) / net_width * seg_width) - rang_x;
  128. int rang_h = ceil(((temp_rect.y + temp_rect.height) * params[1] + params[3]) / net_height * seg_height) - rang_y;
  129. //如果下面的 mask_protos(roi_rangs).clone()位置报错,说明你的output.box数据不对,或者矩形框就1个像素的,开启下面的注释部分防止报错。
  130. rang_w = MAX(rang_w, 1);
  131. rang_h = MAX(rang_h, 1);
  132. if (rang_x + rang_w > seg_width) {
  133. if (seg_width - rang_x > 0)
  134. rang_w = seg_width - rang_x;
  135. else
  136. rang_x -= 1;
  137. }
  138. if (rang_y + rang_h > seg_height) {
  139. if (seg_height - rang_y > 0)
  140. rang_h = seg_height - rang_y;
  141. else
  142. rang_y -= 1;
  143. }
  144. std::vector<cv::Range> roi_rangs;
  145. roi_rangs.push_back(cv::Range(0, 1));
  146. roi_rangs.push_back(cv::Range::all());
  147. roi_rangs.push_back(cv::Range(rang_y, rang_h + rang_y));
  148. roi_rangs.push_back(cv::Range(rang_x, rang_w + rang_x));
  149. //crop
  150. cv::Mat temp_mask_protos = maskProtos(roi_rangs).clone();
  151. cv::Mat protos = temp_mask_protos.reshape(0, { seg_channels,rang_w * rang_h });
  152. cv::Mat matmul_res = (maskProposals * protos).t();
  153. cv::Mat masks_feature = matmul_res.reshape(1, { rang_h,rang_w });
  154. cv::Mat dest, mask;
  155. //sigmoid
  156. cv::exp(-masks_feature, dest);
  157. dest = 1.0 / (1.0 + dest);
  158. int left = floor((net_width / seg_width * rang_x - params[2]) / params[0]);
  159. int top = floor((net_height / seg_height * rang_y - params[3]) / params[1]);
  160. int width = ceil(net_width / seg_width * rang_w / params[0]);
  161. int height = ceil(net_height / seg_height * rang_h / params[1]);
  162. resize(dest, mask, cv::Size(width, height), cv::INTER_NEAREST);
  163. cv::Rect mask_rect = temp_rect - cv::Point(left, top);
  164. mask_rect &= cv::Rect(0, 0, width, height);
  165. mask = mask(mask_rect) > mask_threshold;
  166. if (mask.rows != temp_rect.height || mask.cols != temp_rect.width) { //https://github.com/UNeedCryDear/yolov8-opencv-onnxruntime-cpp/pull/30
  167. resize(mask, mask, temp_rect.size(), cv::INTER_NEAREST);
  168. }
  169. output.boxMask = mask;
  170. }
  171. int BBox2Obb(float centerX, float centerY, float boxW, float boxH, float angle, cv::RotatedRect& rotatedRect) {
  172. rotatedRect = cv::RotatedRect(cv::Point2f(centerX, centerY), cv::Size2f(boxW, boxH), angle);
  173. return 0;
  174. }
  175. void DrawRotatedBox(cv::Mat& srcImg, cv::RotatedRect box, cv::Scalar color, int thinkness) {
  176. cv::Point2f p[4];
  177. box.points(p);
  178. for (int l = 0; l < 4; ++l) {
  179. line(srcImg, p[l], p[(l + 1) % 4], color, thinkness, 8);
  180. }
  181. }
  182. void DrawPred(cv::Mat& img, std::vector<OutputParams> result, std::vector<std::string> classNames, std::vector<cv::Scalar> color) {
  183. cv::Mat mask = img.clone();
  184. for (int i = 0; i < result.size(); i++) {
  185. int left = 0, top = 0;
  186. int color_num = i;
  187. if (result[i].box.area() > 0) {
  188. rectangle(img, result[i].box, color[result[i].id], 2, 8);
  189. left = result[i].box.x;
  190. top = result[i].box.y;
  191. }
  192. if (result[i].rotatedBox.size.width * result[i].rotatedBox.size.height > 0) {
  193. DrawRotatedBox(img, result[i].rotatedBox, color[result[i].id], 2);
  194. left = result[i].rotatedBox.center.x;
  195. top = result[i].rotatedBox.center.y;
  196. }
  197. if (result[i].boxMask.rows && result[i].boxMask.cols > 0)
  198. mask(result[i].box).setTo(color[result[i].id], result[i].boxMask);
  199. std::string label = classNames[result[i].id] + ":" + std::to_string(result[i].confidence);
  200. int baseLine;
  201. cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
  202. top = MAX(top, labelSize.height);
  203. //rectangle(frame, cv::Point(left, top - int(1.5 * labelSize.height)), cv::Point(left + int(1.5 * labelSize.width), top + baseLine), cv::Scalar(0, 255, 0), FILLED);
  204. putText(img, label, cv::Point(left, top), cv::FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
  205. }
  206. cv::addWeighted(img, 0.5, mask, 0.5, 0, img); //add mask to src
  207. }
  208. void DrawPredPose(cv::Mat& img, std::vector<OutputParams> result, PoseParams& poseParams) {
  209. for (int i = 0; i < result.size(); i++) {
  210. int left, top;
  211. int color_num = i;
  212. if (result[i].box.area() > 0) {
  213. rectangle(img, result[i].box, poseParams.objColor, 2, 8);
  214. left = result[i].box.x;
  215. top = result[i].box.y;
  216. }
  217. else
  218. continue;
  219. std::string label = std::to_string(result[i].confidence);
  220. int baseLine;
  221. cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
  222. top = MAX(top, labelSize.height);
  223. //rectangle(frame, cv::Point(left, top - int(1.5 * labelSize.height)), cv::Point(left + int(1.5 * labelSize.width), top + baseLine), cv::Scalar(0, 255, 0), FILLED);
  224. putText(img, label, cv::Point(left+20, top), cv::FONT_HERSHEY_SIMPLEX, 1, poseParams.objColor, 2);
  225. if (result[i].keyPoints.size() != poseParams.kptBodyNames.size())
  226. continue;
  227. for (int j = 0; j < result[i].keyPoints.size(); ++j) {
  228. PoseKeyPoint kpt = result[i].keyPoints[j];
  229. if (kpt.confidence < poseParams.kptThreshold)
  230. continue;
  231. cv::Scalar kptColor = poseParams.posePalette[poseParams.kptColor[j]];
  232. cv::circle(img, cv::Point(kpt.x, kpt.y), poseParams.kptRadius, kptColor, -1, 8);
  233. }
  234. if (poseParams.isDrawKptLine) {
  235. for (int j = 0; j < poseParams.skeleton.size(); ++j) {
  236. PoseKeyPoint kpt0 = result[i].keyPoints[poseParams.skeleton[j][0] - 1];
  237. PoseKeyPoint kpt1 = result[i].keyPoints[poseParams.skeleton[j][1] - 1];
  238. if (kpt0.confidence < poseParams.kptThreshold || kpt1.confidence < poseParams.kptThreshold)
  239. continue;
  240. cv::Scalar kptColor = poseParams.posePalette[poseParams.limbColor[j]];
  241. cv::line(img, cv::Point(kpt0.x, kpt0.y), cv::Point(kpt1.x, kpt1.y), kptColor, 2, 8);
  242. }
  243. }
  244. }
  245. }