Bläddra i källkod

add infer Node

leon 1 månad sedan
förälder
incheckning
4b9227b5af
3 ändrade filer med 66 tillägg och 3 borttagningar
  1. 5 3
      src/main.cpp
  2. 30 0
      src/nodes/infer/inferNode.cpp
  3. 31 0
      src/nodes/infer/inferNodee.hpp

+ 5 - 3
src/main.cpp

@@ -1,10 +1,12 @@
 #include "nodes/base/base.hpp"
 #include "nodes/stream/streamNode.hpp"
-
+#include "nodes/infer/inferNode.hpp"
 int main()
 {
-    std::shared_ptr<Node::BaseNode> node1 = std::make_shared<Node::StreamNode>("test", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/101");
-    node1->start();
+    std::shared_ptr<Node::BaseNode> src_node   = std::make_shared<Node::StreamNode>("src", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/101");
+    std::shared_ptr<Node::BaseNode> infer_node = std::make_shared<Node::InferNode>("infer");
+    infer_node->start();
+    src_node->start();
     while(true)
     {
         std::this_thread::sleep_for(std::chrono::seconds(1));

+ 30 - 0
src/nodes/infer/inferNode.cpp

@@ -0,0 +1,30 @@
+#include "nodes/base/base.hpp"
+#include "nodes/infer/inferNode.hpp"
+#include <unordered_map>
+
+namespace Node
+{
+
+void InferNode::work()
+{
+    printf("InferNode %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;
+            }
+            // do something
+            for (auto& output_buffer : output_buffers_)
+            {
+                output_buffer.second->push(metaData);
+            }
+        }
+    }
+};
+
+}   // namespace Node

+ 31 - 0
src/nodes/infer/inferNodee.hpp

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