#ifndef INFERNODE_HPP__
#define INFERNODE_HPP__

#include "nodes/base/base.hpp"
#include <opencv2/opencv.hpp>
#include "infer/infer.hpp"

namespace Node
{

class InferNode : public BaseNode
{
public:
    InferNode() = delete;
    InferNode(const std::string& name) : BaseNode(name, NODE_TYPE::MID_NODE) {}
    virtual ~InferNode()  { stop(); };

    void set_model_instance(std::shared_ptr<Infer> model, ModelType model_type)
    {
        model_ = model;
        model_type_ = model_type;
    }

    void work() override;

private:
    std::shared_ptr<Infer> model_ = nullptr;
    ModelType model_type_;
};

}

#endif  // INFERNODE_HPP__