videoStreamInfer/src/nodes/stream/streamNode.hpp

#ifndef STREAMNODE_HPP__
#define STREAMNODE_HPP__

#include "nodes/base/base.hpp"
#include "opencv2/opencv.hpp"

#include "stream/stream.hpp"

namespace Node
{

enum class DecodeType
{
    CPU = 0,
    GPU = 1
};

enum class StreamStatus{
    OPENED = 0,
    CLOSED = 1,
    OPEN_FAILED = 2
};

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, int gpu_id=0, DecodeType type=DecodeType::GPU) 
        : BaseNode(name, NODE_TYPE::SRC_NODE), stream_url_(url), gpu_id_(gpu_id), decode_type_(type)
    {
        if (decode_type_ == DecodeType::GPU)
        {
            demuxer_ = FFHDDemuxer::create_ffmpeg_demuxer(stream_url_);
            if (demuxer_ == nullptr)
            {
                printf("demuxer create failed\n");
                status_ = StreamStatus::OPEN_FAILED;
                return;
            }
            decoder_ = FFHDDecoder::create_cuvid_decoder(
                false, FFHDDecoder::ffmpeg2NvCodecId(demuxer_->get_video_codec()), -1, gpu_id, nullptr, nullptr, true
            );
            if (decoder_ == nullptr)
            {
                printf("decoder create failed\n");
                status_ = StreamStatus::OPEN_FAILED;
                return;
            }
            status_ = StreamStatus::OPENED;
        }
        else
        {
            cap_ = std::make_shared<cv::VideoCapture>(stream_url_);
            if (!cap_->isOpened())
            {
                printf("cap open failed\n");
                status_ = StreamStatus::OPEN_FAILED;
                return;
            }
            status_ = StreamStatus::OPENED;
        }
    }
    virtual ~StreamNode() 
    { 
        if (demuxer_) 
        {
            printf("error : demuxer 未被工作线程释放\n");
        }
    };

    void set_stream_url(const std::string& stream_url)
    {
        stream_url_ = stream_url;
    }

    void set_skip_frame(int skip_frame)
    {
        if (skip_frame < 1)
        {
            skip_frame = 1;
        }
        skip_frame_ = skip_frame;
    }



    void work() override;
    void work_cpu();
    void work_gpu();

private:
    std::string stream_url_;
    int skip_frame_ = 1;
    int frame_count_ = -1;
    int gpu_id_ = 0;

    std::shared_ptr<cv::VideoCapture> cap_ = nullptr;
    std::shared_ptr<FFHDDemuxer::FFmpegDemuxer> demuxer_ = nullptr;
    std::shared_ptr<FFHDDecoder::CUVIDDecoder> decoder_ = nullptr;

    DecodeType decode_type_ = DecodeType::GPU;
    StreamStatus status_ = StreamStatus::CLOSED;
};

}

#endif  // STREAMNODE_HPP__