添加画图节点

src/common/meta.hpp

@@ -9,7 +9,9 @@ namespace meta
 
 struct MetaData{
     cv::Mat image;
+    cv::Mat draw_image;
     data::BoxArray boxes;
+    data::BoxArray result;
 };
 
 using MetaDataPtr = std::shared_ptr<MetaData>;

src/main.cpp

@@ -1,22 +1,28 @@
 #include "nodes/base/base.hpp"
 #include "nodes/stream/streamNode.hpp"
 #include "nodes/infer/inferNode.hpp"
+#include "nodes/draw/drawNode.hpp"
+
+
 int main()
 {
     std::shared_ptr<Node::StreamNode> src_node   = std::make_shared<Node::StreamNode>("src", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/101");
     src_node->set_skip_frame(10);
 
 
-    std::shared_ptr<Node::InferNode> infer_node1 = std::make_shared<Node::InferNode>("infer1");
-    std::shared_ptr<Node::InferNode> infer_node2 = std::make_shared<Node::InferNode>("infer2");
-    Node::LinkNode(src_node, infer_node1);
-    Node::LinkNode(infer_node1, infer_node2);
-    infer_node2->start();
-    infer_node1->start();
+    std::shared_ptr<Node::InferNode> infer_node = std::make_shared<Node::InferNode>("infer1");
+    std::shared_ptr<Node::DrawNode> draw_node = std::make_shared<Node::DrawNode>("infer2");
+    Node::LinkNode(src_node, infer_node);
+    Node::LinkNode(infer_node, draw_node);
+    draw_node->start();
+    infer_node->start();
     src_node->start();
+    
     while(true)
     {
         std::this_thread::sleep_for(std::chrono::seconds(1));
     }
     return 0;
-}
+}
+
+// TODO: 单模型多路复用

src/nodes/draw/drawNode.cpp

@@ -0,0 +1,37 @@
+#include "nodes/base/base.hpp"
+#include "nodes/draw/drawNode.hpp"
+namespace Node
+{
+
+void DrawNode::work() 
+{
+    printf("DrawNode %s\n", name_.c_str());
+    while (running_)
+    {
+        for (auto& input_buffer : input_buffers_)
+        {
+            std::shared_ptr<meta::MetaData> metaData;
+            if (!input_buffer.second->try_pop(metaData))
+            {
+                continue;
+            }
+            printf("Node %s get data from %s\n", name_.c_str(), input_buffer.first.c_str());
+            // do something
+            cv::Mat image = metaData->image.clone();
+            for (auto& box : metaData->boxes)
+            {
+                cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(0, 255, 0), 2);
+                cv::putText(image, box.label, cv::Point(box.left, box.top), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0), 2);
+            }
+            metaData->draw_image = image;
+            cv::imwrite("dtest.jpg", image);
+            for (auto& output_buffer : output_buffers_)
+            {
+                printf("Node %s push data to %s\n", name_.c_str(), output_buffer.first.c_str());
+                output_buffer.second->push(metaData);
+            }
+        }
+    }
+}
+
+}

src/nodes/draw/drawNode.hpp

@@ -0,0 +1,27 @@
+#ifndef DRAWNODE_HPP__
+#define DRAWNODE_HPP__
+#include "nodes/base/base.hpp"
+#include <opencv2/opencv.hpp>
+
+namespace Node
+{
+
+class DrawNode : public BaseNode
+{
+public:
+    DrawNode() = delete;
+    DrawNode(const std::string& name) : BaseNode(name, NODE_TYPE::DES_NODE) {}
+    virtual ~DrawNode() { };
+
+    void set_model_path(const std::string& model_path)
+    {
+        model_path_ = model_path;
+    }
+
+    void work() override;
+};
+
+} // namespace Node
+
+
+#endif // DRAWNODE_HPP__

src/nodes/infer/inferNode.cpp

@@ -20,6 +20,11 @@ void InferNode::work()
             }
             printf("Node %s get data from %s\n", name_.c_str(), input_buffer.first.c_str());
             // do something
+            // cv::Mat image = metaData->image;
+
+            // cv::imwrite("test.jpg", image);
+            metaData->boxes.push_back(meta::Box(0, 0, 100, 100, 0.9, "test"));
+            metaData->boxes.push_back(meta::Box(0, 0, 100, 100, 0.9, "test"));
             for (auto& output_buffer : output_buffers_)
             {
                 printf("Node %s push data to %s\n", name_.c_str(), output_buffer.first.c_str());