ElectronPlateDectect.cpp 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
  1. #include "ElectronPlateDectect.h"
  2. #include "JdydAlgorithnm.h"
  3. using namespace std;
  4. //NumberDectect numberDectect;
  5. bool ElectronPlateDectect::Init(bool isCuda)
  6. {
  7. string model_path = "models/electron_plate-sim.onnx";
  8. try {
  9. net = cv::dnn::readNet(model_path);
  10. /*net = cv::dnn::readNet(model_path);
  11. atmosphericPressureALGO.Init();
  12. circularArresterCurrentALGO.Init();*/
  13. }
  14. catch (const std::exception& ex)
  15. {
  16. YunDaISASImageRecognitionService::ConsoleLog(ex.what());
  17. return false;
  18. }
  19. //cuda
  20. if (isCuda) {
  21. net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
  22. net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
  23. }
  24. //cpu
  25. else {
  26. net.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
  27. net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
  28. }
  29. return true;
  30. return false;
  31. }
  32. IDetection::DectectResult ElectronPlateDectect::GetStateResult(cv::Mat img, cv::Rect rec)
  33. {
  34. return resultValue;
  35. }
  36. IDetection::DectectResult ElectronPlateDectect::GetDigitResult(cv::Mat img, cv::Rect rec)
  37. {
  38. //resultValue.clear();
  39. //std::cout << "test" << std::endl;
  40. //try
  41. //{
  42. // cv::Mat ROI = img(rec);
  43. // /*imwrite("test.png", ROI);
  44. // YunDaISASImageRecognitionService::SetImage(QString::fromStdString("test.png"));*/
  45. // Detect(ROI);
  46. //}
  47. //catch (const std::exception& ex)
  48. //{
  49. // YunDaISASImageRecognitionService::ConsoleLog(ex.what());
  50. //}
  51. //if (resultValue.m_confidence < 0.1)
  52. //{
  53. // resultValue = DectectResult(0.45, 0, "");
  54. //}
  55. return resultValue;
  56. }
  57. vector<IDetection::DectectResult> ElectronPlateDectect::GetDigitResults(cv::Mat img, cv::Rect rec)
  58. {
  59. resultValues.clear();
  60. try
  61. {
  62. cv::Mat ROI = img(rec);
  63. Detect(ROI);
  64. if (YunDaISASImageRecognitionService::GetIsShowDectect())
  65. {
  66. cv::Mat drawROI;
  67. ROI.copyTo(drawROI);
  68. DrawPred(drawROI, output, className);
  69. imwrite("test.png", drawROI);
  70. YunDaISASImageRecognitionService::SetImage(QString::fromStdString("test.png"));
  71. }
  72. }
  73. catch (const std::exception& ex)
  74. {
  75. YunDaISASImageRecognitionService::ConsoleLog(ex.what());
  76. }
  77. /*if (resultValue.m_confidence < 0.1)
  78. {
  79. resultValue = DectectResult(0.45, 0, "");
  80. }*/
  81. return resultValues;
  82. }
  83. bool ElectronPlateDectect::Detect(cv::Mat& SrcImg)
  84. {
  85. cv::Mat blob;
  86. int col = SrcImg.cols;
  87. int row = SrcImg.rows;
  88. int maxLen = MAX(col, row);
  89. cv::Mat netInputImg = SrcImg.clone();
  90. if (maxLen > 1.2 * col || maxLen > 1.2 * row) {
  91. cv::Mat resizeImg = cv::Mat::zeros(maxLen, maxLen, CV_8UC3);
  92. SrcImg.copyTo(resizeImg(cv::Rect(0, 0, col, row)));
  93. netInputImg = resizeImg;
  94. }
  95. cv::dnn::blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(0, 0, 0), true, false);
  96. net.setInput(blob);
  97. std::vector<cv::Mat> netOutputImg;
  98. net.forward(netOutputImg, net.getUnconnectedOutLayersNames());
  99. std::vector<int> classIds;//结果id数组
  100. std::vector<float> confidences;//结果每个id对应置信度数组
  101. std::vector<cv::Rect> boxes;//每个id矩形框
  102. float ratio_h = (float)netInputImg.rows / netHeight;
  103. float ratio_w = (float)netInputImg.cols / netWidth;
  104. int net_width = className.size() + 5; //输出的网络宽度是类别数+5
  105. float* pdata = (float*)netOutputImg[0].data;
  106. for (int stride = 0; stride < strideSize; stride++) { //stride
  107. int grid_x = (int)(netWidth / netStride[stride]);
  108. int grid_y = (int)(netHeight / netStride[stride]);
  109. for (int anchor = 0; anchor < 3; anchor++) { //anchors
  110. const float anchor_w = netAnchors[stride][anchor * 2];
  111. const float anchor_h = netAnchors[stride][anchor * 2 + 1];
  112. for (int i = 0; i < grid_y; i++) {
  113. for (int j = 0; j < grid_x; j++) {
  114. float box_score = pdata[4]; ;//获取每一行的box框中含有某个物体的概率
  115. if (box_score >= boxThreshold) {
  116. cv::Mat scores(1, className.size(), CV_32FC1, pdata + 5);
  117. cv::Point classIdPoint;
  118. double max_class_socre;
  119. minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
  120. max_class_socre = (float)max_class_socre;
  121. if (max_class_socre >= classThreshold)
  122. {
  123. //rect [x,y,w,h]
  124. float x = pdata[0]; //x
  125. float y = pdata[1]; //y
  126. float w = pdata[2]; //w
  127. float h = pdata[3]; //h
  128. int left = (x - 0.5 * w) * ratio_w;
  129. int top = (y - 0.5 * h) * ratio_h;
  130. left = left < 0 ? 0 : left;
  131. top = top < 0 ? 0 : top;
  132. int widthBox = int(w * ratio_w);
  133. int heightBox = int(h * ratio_h);
  134. widthBox = widthBox > col ? col : widthBox;
  135. heightBox = heightBox > row ? row : heightBox;
  136. if (left < 0 || left>col || top < 0 || top>row || widthBox > col || heightBox > row
  137. || left + widthBox > col || top + heightBox > row
  138. )
  139. {
  140. continue;
  141. }
  142. classIds.push_back(classIdPoint.x);
  143. confidences.push_back(max_class_socre * box_score);
  144. boxes.push_back(cv::Rect(left, top, widthBox, heightBox));
  145. }
  146. }
  147. pdata += net_width;//下一行
  148. }
  149. }
  150. }
  151. }
  152. //执行非最大抑制以消除具有较低置信度的冗余重叠框(NMS)
  153. vector<int> nms_result;
  154. cv::dnn::NMSBoxes(boxes, confidences, nmsScoreThreshold, nmsThreshold, nms_result);
  155. float confidenceMax = 0;
  156. int confidenceMaxId = 0;
  157. output.clear();
  158. if (nms_result.size() > 0)
  159. {
  160. float confidenceMax = 0;
  161. int confidenceMaxId = -1;
  162. for (int i = 0; i < nms_result.size(); i++) {
  163. int idx = nms_result[i];
  164. if (confidences[idx] > confidenceMax)
  165. {
  166. confidenceMax = confidences[idx];
  167. confidenceMaxId = i;
  168. }
  169. }
  170. if (confidenceMaxId>-1)
  171. {
  172. int idx = nms_result[confidenceMaxId];
  173. Output result(classIds[idx], confidences[idx], boxes[idx]);
  174. output.push_back(result);
  175. if (className[classIds[idx]] == "digit")
  176. {
  177. JdydAlgorithnm::GeNumberDectectInstance().Detect(SrcImg);
  178. if (JdydAlgorithnm::GeNumberDectectInstance().GetDigitResults().size() > 0)
  179. {
  180. for (size_t i = 0; i < JdydAlgorithnm::GeNumberDectectInstance().GetDigitResults().size(); i++)
  181. {
  182. auto tempResultValue = DectectResult(confidences[idx],
  183. JdydAlgorithnm::GeNumberDectectInstance().GetDigitResults()[i].m_dValue,
  184. className[classIds[idx]]);
  185. if (JdydAlgorithnm::GeNumberDectectInstance().GetDigitResults()[i].m_sValue != "")
  186. {
  187. tempResultValue.m_sValue = className[classIds[idx]] + "_" + JdydAlgorithnm::GeNumberDectectInstance().GetDigitResults()[i].m_sValue;
  188. }
  189. resultValues.push_back(tempResultValue);
  190. }
  191. }
  192. }
  193. else if (className[classIds[idx]] == "state")
  194. {
  195. auto tempResultValue = DectectResult(confidences[idx], 0, className[classIds[idx]] + "_green_on");
  196. resultValues.push_back(tempResultValue);
  197. }
  198. YunDaISASImageRecognitionService::ConsoleLog(QString::fromStdString(className[classIds[idx]]));
  199. }
  200. }
  201. else {
  202. resultValue = DectectResult(confidenceMax, 0.00, "");
  203. }
  204. return false;
  205. }