ClassifyDectect.h 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. #include<iostream>
  2. #include <opencv2/opencv.hpp>
  3. #include <opencv2/highgui/highgui_c.h>
  4. #include <iostream>
  5. #include <vector>
  6. #include <algorithm>
  7. #include <cstdlib>
  8. #include <fstream>
  9. //#include "MeterRead.h"
  10. using namespace std;
  11. //using namespace cv::dnn;
  12. #define YOLO_P6 false //是否使用P6模型
  13. struct Output {
  14. int id; //结果类别id
  15. float confidence; //结果置信度
  16. cv::Rect box; //矩形框
  17. };
  18. class ClassifyDectect {
  19. public:
  20. ClassifyDectect() {
  21. }
  22. ~ClassifyDectect() {}
  23. bool Init(bool isCuda);
  24. bool Detect(cv::Mat& SrcImg, std::vector<Output>& output);
  25. void drawPred(cv::Mat& img, std::vector<Output> result);
  26. void modifyConfidenceParameter(float boxThreshold, float classThreshold, float nmsThreshold);
  27. private:
  28. cv::dnn::Net net;
  29. #if(defined YOLO_P6 && YOLO_P6==true)
  30. const float netAnchors[4][6] = { { 19,27, 44,40, 38,94 },{ 96,68, 86,152, 180,137 },{ 140,301, 303,264, 238,542 },{ 436,615, 739,380, 925,792 } };
  31. const int netWidth = 1280; //ONNX图片输入宽度
  32. const int netHeight = 1280; //ONNX图片输入高度
  33. const int strideSize = 4; //stride size
  34. #else
  35. const float netAnchors[3][6] = { { 10,13, 16,30, 33,23 },{ 30,61, 62,45, 59,119 },{ 116,90, 156,198, 373,326 } };
  36. const int netWidth = 640; //ONNX图片输入宽度
  37. const int netHeight = 640; //ONNX图片输入高度
  38. const int strideSize = 3; //stride size
  39. #endif // YOLO_P6
  40. const float netStride[4] = { 8, 16.0,32,64 };
  41. float boxThreshold = 0.35;
  42. float classThreshold = 0.35;
  43. float nmsThreshold = 0.55;
  44. float nmsScoreThreshold = boxThreshold * classThreshold;
  45. public:
  46. const std::vector<std::string> className = {
  47. "disconnector",
  48. "oil_level",
  49. "meter",
  50. "changeover_switch",
  51. "pressplate",
  52. "disconnector_state",
  53. "light",
  54. "air_switch",
  55. "character",
  56. "indoor_disconnector",
  57. "triangle_disconnector",
  58. "electron_plate",
  59. "contact",
  60. "instructions",
  61. "operating_handle",
  62. "person",
  63. "fire",
  64. "smoke",
  65. "circle_off_on_plate"
  66. };
  67. const std::string instructionsName = "instructions";
  68. };