Browse Source

add record node

leon 5 months ago
parent
commit
378340da54
3 changed files with 87 additions and 4 deletions
  1. 7 4
      src/main.cpp
  2. 45 0
      src/nodes/record/recordNode.cpp
  3. 35 0
      src/nodes/record/recordNode.hpp

+ 7 - 4
src/main.cpp

@@ -4,7 +4,7 @@
 #include "nodes/draw/drawNode.hpp"
 #include "nodes/draw/drawNode.hpp"
 #include "nodes/track/trackNode.hpp"
 #include "nodes/track/trackNode.hpp"
 #include "nodes/httpPush/httpPush.hpp"
 #include "nodes/httpPush/httpPush.hpp"
-
+#include "nodes/record/recordNode.hpp"
 
 
 void test_depth()
 void test_depth()
 {
 {
@@ -41,13 +41,16 @@ void test_yolo()
     std::shared_ptr<GNode::TrackNode> track_node     = std::make_shared<GNode::TrackNode>("tracker", 30, 30);
     std::shared_ptr<GNode::TrackNode> track_node     = std::make_shared<GNode::TrackNode>("tracker", 30, 30);
 
 
     std::shared_ptr<GNode::DrawNode> draw_node     = std::make_shared<GNode::DrawNode>("draw");
     std::shared_ptr<GNode::DrawNode> draw_node     = std::make_shared<GNode::DrawNode>("draw");
-    std::shared_ptr<GNode::HttpPushNode> push_node = std::make_shared<GNode::HttpPushNode>("push", "172.16.20.168", 8080, "/push");
+    std::shared_ptr<GNode::RecordNode> record_node = std::make_shared<GNode::RecordNode>("record");
+    record_node->set_record_path("result/");
+    record_node->set_fps(25);
+    record_node->set_fourcc(cv::VideoWriter::fourcc('M', 'J', 'P', 'G'));
     
     
     GNode::LinkNode(src_node0, infer_node);
     GNode::LinkNode(src_node0, infer_node);
     GNode::LinkNode(infer_node, track_node);
     GNode::LinkNode(infer_node, track_node);
     GNode::LinkNode(track_node, draw_node);
     GNode::LinkNode(track_node, draw_node);
-    GNode::LinkNode(draw_node, push_node);
-    push_node->start();
+    GNode::LinkNode(draw_node, record_node);
+    record_node->start();
     draw_node->start();
     draw_node->start();
     track_node->start();
     track_node->start();
     infer_node->start();
     infer_node->start();

+ 45 - 0
src/nodes/record/recordNode.cpp

@@ -0,0 +1,45 @@
+#include "common/data.hpp"
+#include "common/meta.hpp"
+#include "nodes/recorder/recorderNode.hpp"
+
+namespace GNode
+{
+
+void RecordNode::work()
+{
+    printf("RecordNode %s\n", name_.c_str());
+    while (running_)
+    {
+        bool has_data = false;
+        for (auto& input_buffer : input_buffers_)
+        {
+            std::shared_ptr<meta::MetaData> metaData;
+            if (!input_buffer.second->try_pop(metaData))
+            {
+                continue;
+            }
+            has_data = true;
+            // printf("Node %s get data from %s\n", name_.c_str(), input_buffer.first.c_str());
+            cv::Mat image = metaData->draw_image;
+            if (!writer_.isOpened())
+            {
+                writer_.open(record_path_, fourcc, fps_, cv::Size(image.cols, image.rows));
+                if (!writer_.isOpened())
+                {
+                    std::cerr << "Error: cannot open video writer" << std::endl;
+                    break;
+                }
+            }
+            writer_ << image;
+        }
+        if (!has_data) {
+            std::unique_lock<std::mutex> lock(mutex_);
+            cond_var_->wait_for(lock, std::chrono::milliseconds(100), [this] {
+                return !running_;  // 等待时检查退出条件
+            });
+        }
+    }
+}
+
+
+}   // namespace GNode

+ 35 - 0
src/nodes/record/recordNode.hpp

@@ -0,0 +1,35 @@
+#ifndef RECORDNODE_HPP__
+#define RECORDNODE_HPP__
+#include "nodes/base/base.hpp"
+#include <opencv2/opencv.hpp>
+
+namespace GNode
+{
+
+class RecordNode : public BaseNode
+{
+public:
+    RecordNode() = delete;
+    RecordNode(const std::string& name) : BaseNode(name, NODE_TYPE::DES_NODE) {}
+    virtual ~RecordNode() 
+    { 
+        stop(); 
+        if (writer_.isOpened())
+        {
+            writer_.release();
+        }
+    }
+
+    void work() override;
+    void set_record_path(const std::string& path) { record_path_ = path; }
+    void set_fps(int fps) { fps_ = fps; }
+    void set_fourcc(cv::VideoWriter::fourcc) { fourcc_ = fourcc; }
+private:
+    std::string record_path_;
+    int fps_ = 25;
+    cv::VideoWriter::fourcc_ = cv::VideoWriter::fourcc('M', 'J', 'P', 'G');
+    cv::VideoWriter writer_;
+};
+    
+}
+#endif // RECORDNODE_HPP__