Bläddra i källkod

初始化仓库

leon 1 månad sedan
incheckning
14de4d5977
6 ändrade filer med 371 tillägg och 0 borttagningar
  1. 74 0
      src/common/data.hpp
  2. 24 0
      src/common/meta.hpp
  3. 59 0
      src/common/queue.hpp
  4. 49 0
      src/nodes/base/base.cpp
  5. 101 0
      src/nodes/base/base.hpp
  6. 64 0
      src/nodes/stream/streamNode.hpp

+ 74 - 0
src/common/data.hpp

@@ -0,0 +1,74 @@
+#ifndef DATA_HPP__
+#define DATA_HPP__
+
+#include <string>
+#include <vector>
+#include <map>
+#include <iostream>
+
+
+
+namespace data
+{
+
+struct Point
+{
+    float x;
+    float y;
+    float score;
+
+    Point() : x(0), y(0), score(0) {}
+    Point(float x, float y, float score) : x(x), y(y), score(score) {}
+    Point(const Point& p) : x(p.x), y(p.y), score(p.score) {}
+    Point& operator=(const Point& p)
+    {
+        x = p.x;
+        y = p.y;
+        score = p.score;
+        return *this;
+    }
+};
+
+
+struct Segment
+{
+
+};
+
+struct Box
+{
+    float left;
+    float top;
+    float right;
+    float bottom;
+    float score;
+    std::string label;
+    // 是否需要二次分类
+    bool need_classification;
+    std::vector<Point> points;
+    Box() : left(0), top(0), right(0), bottom(0), score(0), label(""), need_classification(false) {}
+    Box(float left, float top, float right, float bottom, float score, const std::string& label, bool need_classification = false) 
+        : left(left), top(top), right(right), bottom(bottom), score(score), label(label), need_classification(need_classification) {}
+    Box(const Box& b) : left(b.left), top(b.top), right(b.right), bottom(b.bottom), score(b.score), label(b.label), need_classification(b.need_classification), points(b.points) {}
+    Box& operator=(const Box& b)
+    {
+        left = b.left;
+        top = b.top;
+        right = b.right;
+        bottom = b.bottom;
+        score = b.score;
+        label = b.label;
+        need_classification = b.need_classification;
+        points = b.points;
+        return *this;
+    }
+};
+
+using BoxArray = std::vector<Box>;
+
+
+
+} // namespace data
+
+
+#endif // DATA_HPP__

+ 24 - 0
src/common/meta.hpp

@@ -0,0 +1,24 @@
+#ifndef META_HPP__
+#define META_HPP__
+
+#include "common/data.hpp"
+#include <opencv2/opencv.hpp>
+
+namespace meta
+{
+
+struct MetaData{
+    cv::Mat image;
+    data::BoxArray boxes;
+};
+
+using MetaDataPtr = std::shared_ptr<MetaData>;
+
+inline MetaDataPtr makeMetaData(const MetaData& metaData) {
+    return std::make_shared<MetaData>(metaData);
+}
+
+}
+
+
+#endif // META_HPP__

+ 59 - 0
src/common/queue.hpp

@@ -0,0 +1,59 @@
+#ifndef QUEUE_HPP__
+#define QUEUE_HPP__
+
+#include <queue>
+#include <mutex>
+#include <condition_variable>
+
+template <typename T>
+class SharedQueue {
+public:
+    SharedQueue() = default;
+
+    void push(const T& item) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        queue_.push(item);
+        cond_var_.notify_one();
+    }
+
+    void push(T&& item) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        queue_.push(std::move(item));
+        cond_var_.notify_one();
+    }
+
+    bool tryPop(T& item) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        if (queue_.empty()) {
+            return false;
+        }
+        item = std::move(queue_.front());
+        queue_.pop();
+        return true;
+    }
+
+    T waitAndPop() {
+        std::unique_lock<std::mutex> lock(mutex_);
+        cond_var_.wait(lock, [this] { return !queue_.empty(); });
+        T item = std::move(queue_.front());
+        queue_.pop();
+        return item;
+    }
+
+    bool empty() const {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return queue_.empty();
+    }
+
+    size_t size() const {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return queue_.size();
+    }
+
+private:
+    mutable std::mutex mutex_;
+    std::queue<T> queue_;
+    std::condition_variable cond_var_;
+};
+
+#endif  // QUEUE_HPP__

+ 49 - 0
src/nodes/base/base.cpp

@@ -0,0 +1,49 @@
+#include "common/meta.hpp"
+#include "nodes/base/base.hpp"
+
+namespace Node
+{
+
+BaseNode::BaseNode(const std::string& name, NODE_TYPE type)
+{
+    name_ = name;
+    type_ = type;
+}
+
+BaseNode::~BaseNode()
+{
+    stop();
+}
+
+void BaseNode::start()
+{
+    if (!running_)
+    {
+        worker_thread_ = std::thread(&BaseNode::work, this);
+        running_ = true;
+    }
+    else
+    {
+        std::cerr << "Error: node is already running" << std::endl;
+    }
+}
+
+void BaseNode::stop()
+{
+    if (running_)
+    {
+        cond_var_->notify_all();
+        running_ = false;
+        if (worker_thread_.joinable())
+        {
+            worker_thread_.join();
+        }
+        printf("Node %s stopped\n", name_.c_str());
+    }
+    else
+    {
+        printf("Node %s already running\n", name_.c_str());
+    }
+}
+
+} // namespace Node

+ 101 - 0
src/nodes/base/base.hpp

@@ -0,0 +1,101 @@
+#ifndef BASE_HPP__
+#define BASE_HPP__
+
+#include "common/meta.hpp"
+#include "common/queue.hpp"
+#include "common/data.hpp"
+#include <queue>
+#include <unordered_map>
+#include <string>
+#include <thread>
+#include <mutex>
+#include <condition_variable>
+#include <memory>
+
+namespace Node
+{
+
+enum NODE_TYPE {
+    SRC_NODE,  // 输入节点
+    MID_NODE,  // 中间节点
+    DES_NODE,  // 输出节点
+};
+
+class BaseNode
+{
+public:
+    BaseNode() = delete;
+    BaseNode(const std::string& name, NODE_TYPE type);
+    virtual ~BaseNode();
+
+    virtual void work() = 0;
+
+    void start();
+    void stop();
+
+    inline void add_input_buffer(const std::string& name, SharedQueue buffer)
+    {
+        std::unique_lock<std::mutex> lock(mutex_);
+        input_buffers_[name] = buffer;
+    }
+
+    inline void add_output_buffer(const std::string& name, SharedQueue buffer)
+    {
+        std::unique_lock<std::mutex> lock(mutex_);
+        output_buffers_[name] = buffer;
+    }
+
+    inline void del_input_buffer(const std::string& name)
+    {
+        std::unique_lock<std::mutex> lock(mutex_);
+        if (input_buffers_.find(name) == input_buffers_.end())
+        {
+            return;
+        }
+        input_buffers_.erase(name);
+    }
+
+    inline void del_output_buffer(const std::string& name)
+    {
+        std::unique_lock<std::mutex> lock(mutex_);
+        if (output_buffers_.find(name) == output_buffers_.end())
+        {
+            return;
+        }
+        output_buffers_.erase(name);
+    }
+
+    inline void is_running()
+    {
+        return running_;
+    }
+
+    inline std::string get_name()
+    {
+        return name_;
+    }
+
+protected:
+    std::string name_;
+    NODE_TYPE type_;
+    std::thread worker_thread_;
+    std::mutex mutex_;
+    std::shared_ptr<std::condition_variable> cond_var_ =
+        std::make_shared<std::condition_variable>();
+    bool running_ = false;
+    std::unordered_map<std::string, SharedQueue<std::shared_ptr<meta::MetaData>>> input_buffers_;
+    std::unordered_map<std::string, SharedQueue<std::shared_ptr<meta::MetaData>>> output_buffers_;
+};
+
+static inline void LinkNode(const std::shared_ptr<BaseNode> &front,
+    const std::shared_ptr<BaseNode> &back) 
+{
+    auto queue = std::make_shared<SharedQueue>();
+    back->add_input(front->get_name(), queue);
+    front->add_output(back->get_name(), queue);
+}
+
+
+}
+
+#endif // BASE_HPP__

+ 64 - 0
src/nodes/stream/streamNode.hpp

@@ -0,0 +1,64 @@
+#ifndef STREAMNODE_HPP__
+#define STREAMNODE_HPP__
+
+#include "nodes/base/base.hpp"
+#include <opencv2/opencv.hpp>
+
+namespace Node
+{
+
+class StreamNode : public BaseNode
+{
+public:
+    StreamNode() = delete;
+    StreamNode(const std::string& name) : BaseNode(name, NODE_TYPE::SRC_NODE) {}
+    StreamNode(const std::string& name, const std::string& url) : BaseNode(name, NODE_TYPE::SRC_NODE), stream_url_(url){}
+    virtual ~StreamNode() { };
+
+    void set_stream_url(const std::string& stream_url)
+    {
+        stream_url_ = stream_url;
+    }
+
+    void set_skip_frame(int skip_frame)
+    {
+        skip_frame_ = skip_frame;
+    }
+
+    void work() override
+    {
+        cv::VideoCapture cap(stream_url_);
+        if (!cap.isOpened())
+        {
+            std::cerr << "Error: cannot open camera" << std::endl;
+            return;
+        }
+
+        while (running_)
+        {
+            cv::Mat frame;
+            cap >> frame;
+            if (frame.empty())
+            {
+                std::cerr << "Error: cannot read frame" << std::endl;
+                break;
+            }
+
+            auto metaData = std::make_shared<meta::MetaData>();
+            metaData->image = frame;
+
+            for (auto& output_buffer : output_buffers_)
+            {
+                output_buffer.second.push(metaData);
+            }
+        }
+    }
+
+private:
+    std::string stream_url_;
+    int skip_frame_ = 0;
+};
+
+}
+
+#endif  // STREAMNODE_HPP__