main.cpp 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. #include "nodes/base/base.hpp"
  2. #include "nodes/stream/streamNode.hpp"
  3. #include "nodes/infer/inferNode.hpp"
  4. #include "nodes/draw/drawNode.hpp"
  5. #include "nodes/track/trackNode.hpp"
  6. #include "nodes/httpPush/httpPush.hpp"
  7. #include "nodes/record/recordNode.hpp"
  8. void test_depth()
  9. {
  10. std::shared_ptr<GNode::StreamNode> src_node0 = std::make_shared<GNode::StreamNode>("src0", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/201", 0, GNode::DecodeType::GPU);
  11. src_node0->set_skip_frame(1);
  12. std::shared_ptr<Infer> depth_model = load("model/depth.engine", ModelType::DEPTH_ANYTHING, {}, 0, 0.25, 0.45);
  13. std::shared_ptr<GNode::InferNode> infer_node = std::make_shared<GNode::InferNode>("depth");
  14. infer_node->set_model_instance(depth_model, ModelType::DEPTH_ANYTHING);
  15. std::shared_ptr<GNode::DrawNode> draw_node = std::make_shared<GNode::DrawNode>("draw");
  16. std::shared_ptr<GNode::HttpPushNode> push_node = std::make_shared<GNode::HttpPushNode>("push", "172.16.20.168", 8080, "/push");
  17. GNode::LinkNode(src_node0, infer_node);
  18. GNode::LinkNode(infer_node, draw_node);
  19. GNode::LinkNode(draw_node, push_node);
  20. push_node->start();
  21. draw_node->start();
  22. infer_node->start();
  23. src_node0->start();
  24. std::this_thread::sleep_for(std::chrono::seconds(10));
  25. }
  26. void test_yolo()
  27. {
  28. // std::vector<std::string> names = { "person", "clothes", "vest" };
  29. std::vector<std::string> names = { "person", "car", "close", "open" };
  30. // std::shared_ptr<GNode::StreamNode> src_node0 = std::make_shared<GNode::StreamNode>("src0", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/201", 0, GNode::DecodeType::GPU);
  31. std::shared_ptr<GNode::StreamNode> src_node0 = std::make_shared<GNode::StreamNode>("src0", "carperson.mp4", 0, GNode::DecodeType::GPU);
  32. src_node0->set_skip_frame(1);
  33. std::shared_ptr<Infer> yolo_model = load("model/carperson.engine", ModelType::YOLOV5, names, 0, 0.25, 0.45);
  34. std::shared_ptr<GNode::InferNode> infer_node = std::make_shared<GNode::InferNode>("yolov5");
  35. infer_node->set_model_instance(yolo_model, ModelType::YOLO11);
  36. std::shared_ptr<GNode::TrackNode> track_node = std::make_shared<GNode::TrackNode>("tracker", "person", 30, 30);
  37. std::shared_ptr<GNode::DrawNode> draw_node = std::make_shared<GNode::DrawNode>("draw_track");
  38. std::shared_ptr<GNode::RecordNode> record_node = std::make_shared<GNode::RecordNode>("record");
  39. record_node->set_record_path("result/result.mp4");
  40. record_node->set_fps(25);
  41. record_node->set_fourcc(cv::VideoWriter::fourcc('X', '2', '6', '4'));
  42. GNode::LinkNode(src_node0, infer_node);
  43. // GNode::LinkNode(infer_node, track_node);
  44. GNode::LinkNode(infer_node, draw_node);
  45. // GNode::LinkNode(draw_node, record_node);
  46. // record_node->start();
  47. draw_node->start();
  48. // track_node->start();
  49. infer_node->start();
  50. src_node0->start();
  51. getchar();
  52. }
  53. void test_multi()
  54. {
  55. std::vector<std::string> names = { "person", "clothes", "vest" };
  56. std::shared_ptr<GNode::StreamNode> src_node0 = std::make_shared<GNode::StreamNode>("src0", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/101", 0, GNode::DecodeType::GPU);
  57. src_node0->set_skip_frame(10);
  58. std::shared_ptr<GNode::StreamNode> src_node1 = std::make_shared<GNode::StreamNode>("src1", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/201", 0, GNode::DecodeType::GPU);
  59. src_node1->set_skip_frame(10);
  60. std::shared_ptr<GNode::StreamNode> src_node2 = std::make_shared<GNode::StreamNode>("src2", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/301", 0, GNode::DecodeType::GPU);
  61. src_node2->set_skip_frame(10);
  62. std::shared_ptr<Infer> yolo_model = load("model/yolo11s.engine", ModelType::YOLO11, names, 0, 0.25, 0.45);
  63. std::shared_ptr<GNode::InferNode> infer_node1 = std::make_shared<GNode::InferNode>("yolo11");
  64. infer_node1->set_model_instance(yolo_model, ModelType::YOLO11);
  65. std::shared_ptr<Infer> depth_model = load("model/depth.engine", ModelType::DEPTH_ANYTHING, {}, 0, 0.25, 0.45);
  66. std::shared_ptr<GNode::InferNode> infer_node2 = std::make_shared<GNode::InferNode>("depth");
  67. infer_node2->set_model_instance(depth_model, ModelType::DEPTH_ANYTHING);
  68. std::shared_ptr<GNode::DrawNode> draw_node = std::make_shared<GNode::DrawNode>("draw");
  69. std::shared_ptr<GNode::HttpPushNode> push_node = std::make_shared<GNode::HttpPushNode>("push", "172.16.20.168", 8080, "/push");
  70. GNode::LinkNode(src_node0, infer_node1);
  71. GNode::LinkNode(src_node1, infer_node1);
  72. GNode::LinkNode(src_node2, infer_node1);
  73. GNode::LinkNode(infer_node1, infer_node2);
  74. GNode::LinkNode(infer_node2, draw_node);
  75. GNode::LinkNode(draw_node, push_node);
  76. push_node->start();
  77. draw_node->start();
  78. infer_node1->start();
  79. infer_node2->start();
  80. src_node0->start();
  81. src_node1->start();
  82. src_node2->start();
  83. getchar();
  84. }
  85. int main()
  86. {
  87. test_yolo();
  88. return 0;
  89. }
  90. // TODO
  91. // 硬解码、软解码 完成 软解使用opencv + ffmpeg,硬解使用nvcodec + ffmpeg
  92. // 模型复用 完成,基类加锁保证一致性
  93. // 画图节点 完成
  94. // 推送节点 基本完成
  95. // 分析节点
  96. // depth-anything 进行中
  97. // YOLO11 seg
  98. // 通过配置文件创建 pipeline
  99. // 日志
  100. // 设置电子围栏