利用ffmpeg 从RTSP流 取rgb图

一，思路

拉rtsp视频流-》解析出yuv frame->转成 rgb  mat

子线程 负责解析 yuc frame，转rgb mat 并存入一个队列

主线程一共一个 取图的接口，调用时从队列取出一个mat

二，代码

#pragma once

#include <stdio.h>
#include <istream>
#include <iostream>
#include <inttypes.h>
#include <string>
#include <stdlib.h>
#include <unistd.h>
#include <atomic>
#include <thread>
#include <grab/stream/lockfree_queue.h>

// ffmpeg
#ifdef __cplusplus
extern "C"
{
#endif

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
#include <libswresample/swresample.h>
#include <libavdevice/avdevice.h>
#include <libavformat/version.h>
#include <libavutil/mathematics.h>
#include <libavutil/imgutils.h>
#include <libavutil/time.h>
#include <libavutil/frame.h>

#ifdef __cplusplus
}
#endif

// opencv
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>

namespace Stream
{
    class MediaDecode
    {
    public:
        MediaDecode(const std::string &_rtspUrl);
        ~MediaDecode();

        //初始化
        bool init();
        //开始取流
        void startGrabbing();
        //停止取流
        void stopGrabbing();

        //取图  rgb格式
        cv::Mat grabImage();

    private:
        // rtsp地址
        std::string rtspUrl;
        //初始化结果
        bool isInit{false};
        //是否停止
        std::atomic<bool> isStop{false};
        //frame队列
        LockFreeQueue<cv::Mat> pAVFrameQueue{3};

        //帧率
        double fps = 0;
        //帧间隔
        double interval = 0;

        //ffmpeg的全局上下文
        AVFormatContext *pAVFormatContext = 0;
        //ffmpeg流信息
        AVStream *pAVStream = 0;
        //ffmpeg编码上下文
        AVCodecContext *pAVCodecContext = 0;
        //ffmpeg编码器
        AVCodec *pAVCodec = 0;
        //ffmpag单帧数据包
        AVPacket *pAVPacket = 0;
        //ffmpeg单帧缓存
        AVFrame *pAVFrame = 0;
        //ffmpeg数据字典，用于配置一些编码器属性等
        AVDictionary *pAVDictionary = 0;

        //函数执行结果
        int ret = 0;
        //视频流所在的序号
        int videoIndex = -1;

    private:
        // yuv frame换rgb mat
        cv::Mat avFrameToMat(AVFrame *avframe, int width, int height);
    };
}

#include <grab/stream/media_decode.h>

namespace Stream
{
    MediaDecode::MediaDecode(const std::string &_rtspUrl) : rtspUrl(_rtspUrl)
    {
    }

    MediaDecode::~MediaDecode()
    {
        isStop = true;
        if (pAVFrame)
        {
            av_frame_free(&pAVFrame);
            pAVFrame = 0;
        }
        if (pAVPacket)
        {
            av_free_packet(pAVPacket);
            pAVPacket = 0;
        }
        if (pAVCodecContext)
        {
            avcodec_close(pAVCodecContext);
            pAVCodecContext = 0;
        }
        if (pAVFormatContext)
        {
            avformat_close_input(&pAVFormatContext);
            avformat_free_context(pAVFormatContext);
            pAVFormatContext = 0;
        }
    }

    bool MediaDecode::init()
    {
        //分配
        pAVFormatContext = avformat_alloc_context();
        pAVPacket = av_packet_alloc();
        pAVFrame = av_frame_alloc();

        if (!pAVFormatContext || !pAVPacket || !pAVFrame)
        {
            std::cout << "failed to alloc \n";
            return false;
        }

        //注册所有容器和编解码器
        av_register_all();
        avformat_network_init();

        //打开视频流
        ret = avformat_open_input(&pAVFormatContext, rtspUrl.c_str(), 0, 0);
        if (ret)
        {
            std::cout << "failed to open rtsp \n";
            return false;
        }

        //探测流媒体信息
        ret = avformat_find_stream_info(pAVFormatContext, 0);
        if (ret < 0)
        {
            std::cout << "failed to find stream \n";
            return false;
        }

        //提取视频信息
        for (int index = 0; index < pAVFormatContext->nb_streams; index++)
        {
            pAVCodecContext = pAVFormatContext->streams[index]->codec;
            pAVStream = pAVFormatContext->streams[index];
            switch (pAVCodecContext->codec_type)
            {
            case AVMEDIA_TYPE_UNKNOWN:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_UNKNOWN \n";
                break;
            case AVMEDIA_TYPE_VIDEO:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_VIDEO \n";
                videoIndex = index;
                break;
            case AVMEDIA_TYPE_AUDIO:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_AUDIO \n";
                break;
            case AVMEDIA_TYPE_DATA:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_DATA \n";
                break;
            case AVMEDIA_TYPE_SUBTITLE:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_SUBTITLE \n";
                break;
            case AVMEDIA_TYPE_ATTACHMENT:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_ATTACHMENT \n";
                break;
            case AVMEDIA_TYPE_NB:
                std::cout << "流序号:" << index << "类型为:"
                          << "AVMEDIA_TYPE_NB \n";
                break;
            default:
                break;
            }
            // 已经找打视频品流
            if (videoIndex != -1)
            {
                break;
            }
        }

        if (videoIndex == -1 || !pAVCodecContext)
        {
            std::cout << "failed to find video stream \n";
            return false;
        }

        //找到解码器
        // pAVCodec = avcodec_find_decoder_by_name("h264_nvmpi");
        pAVCodec = avcodec_find_decoder_by_name("h264");
        if (!pAVCodec)
        {
            std::cout << "fialed to find decoder \n";
            return false;
        }

        //打开解码器
        // 设置缓存大小 1024000byte
        av_dict_set(&pAVDictionary, "buffer_size", "1024000", 0);
        // 设置超时时间 20s
        av_dict_set(&pAVDictionary, "stimeout", "20000000", 0);
        // 设置最大延时 3s
        av_dict_set(&pAVDictionary, "max_delay", "30000000", 0);
        // 设置打开方式 tcp/udp
        av_dict_set(&pAVDictionary, "rtsp_transport", "tcp", 0);
        ret = avcodec_open2(pAVCodecContext, pAVCodec, &pAVDictionary);
        if (ret)
        {
            std::cout << "failed to open codec \n";
            return false;
        }

        // 显示视频相关的参数信息（编码上下文）
        std::cout << "比特率:" << pAVCodecContext->bit_rate << std::endl;
        std::cout << "宽高:" << pAVCodecContext->width << "x" << pAVCodecContext->height << std::endl;
        std::cout << "格式:" << pAVCodecContext->pix_fmt << std::endl; // AV_PIX_FMT_YUV420P 0
        std::cout << "帧率分母:" << pAVCodecContext->time_base.den << std::endl;
        std::cout << "帧率分子:" << pAVCodecContext->time_base.num << std::endl;
        std::cout << "帧率分母:" << pAVStream->avg_frame_rate.den << std::endl;
        std::cout << "帧率分子:" << pAVStream->avg_frame_rate.num << std::endl;
        std::cout << "总时长:" << pAVStream->duration / 10000.0 << "s" << std::endl;
        std::cout << "总帧数:" << pAVStream->nb_frames << std::endl;

        fps = pAVStream->avg_frame_rate.num * 1.0f / pAVStream->avg_frame_rate.den;
        interval = 1 * 1000 / fps;
        std::cout << "平均帧率:" << fps << std::endl;
        std::cout << "帧间隔:" << interval << "ms" << std::endl;

        isInit = true;
        return isInit;
    }

    void MediaDecode::startGrabbing()
    {
        if (!isInit)
        {
            return;
        }

        std::thread t([&]()
                      {
                    while (!isStop)
        {
        
        while (av_read_frame(pAVFormatContext, pAVPacket) >= 0)
        {
            if (pAVPacket->stream_index == videoIndex)
            {
                //对读取的数据包进行解码
                ret = avcodec_send_packet(pAVCodecContext, pAVPacket);
                if (ret)
                {
                    std::cout << "failed to avcodec_send_packet(pAVCodecContext, pAVPacket) ,ret =" << ret << std::endl;
                    break;
                }
                while (!avcodec_receive_frame(pAVCodecContext, pAVFrame))
                {
                    //存入队列
                    cv::Mat mat = avFrameToMat(pAVFrame, pAVFrame->width, pAVFrame->height).clone();
                    pAVFrameQueue.push_without_wait(mat);
                }
            }
        }
        } });
        t.detach();
    }

    void MediaDecode::stopGrabbing()
    {
        isStop = true;
    }

    cv::Mat MediaDecode::grabImage()
    {
        cv::Mat mat;
        pAVFrameQueue.pop_without_wait(mat);
        return mat;
    }

    cv::Mat MediaDecode::avFrameToMat(AVFrame *avframe, int width, int height)
    {
        if (width <= 0)
            width = avframe->width;
        if (height <= 0)
            height = avframe->height;
        struct SwsContext *sws_ctx = NULL;
        sws_ctx = sws_getContext(avframe->width, avframe->height, (enum AVPixelFormat)avframe->format, width, height, AV_PIX_FMT_BGR24, SWS_BICUBIC, NULL, NULL, NULL);
        cv::Mat mat;
        mat.create(cv::Size(width, height), CV_8UC3);
        AVFrame *bgr24frame = av_frame_alloc();
        bgr24frame->data[0] = (uint8_t *)mat.data;
        avpicture_fill((AVPicture *)bgr24frame, bgr24frame->data[0], AV_PIX_FMT_BGR24, width, height);
        sws_scale(sws_ctx, (const uint8_t *const *)avframe->data, avframe->linesize, 0, avframe->height, // from cols=0,all rows trans
                  bgr24frame->data, bgr24frame->linesize);
        av_free(bgr24frame);
        sws_freeContext(sws_ctx);
        return mat;
    }
}

三，测试

TEST(stream, stream_test)
{
    MediaDecode mediaDecode("rtsp://xxxxx");
    if (mediaDecode.init())
    {
        mediaDecode.startGrabbing();
        int index = 0;
        while (1)
        {
            this_thread::sleep_for(milliseconds(40));
            cv::Mat rgbMat = mediaDecode.grabImage();
            if (rgbMat.cols>0)
            {
                std::cout << "grabImage" << rgbMat.cols <<"------------------------"<<rgbMat.rows << std::endl;
                char filename[50];
                index++;
                sprintf(filename, "%d.jpg", index);
                imwrite(filename, rgbMat);
            }
        }

        mediaDecode.stopGrabbing();
    }
}

#2022/12/07

将队列改成阻塞，这样再不主动取图时，会停止没必要的抽帧。

/*
 * @Author: keiler [email protected]
 * @Date: 2022-12-07 09:33:27
 * @LastEditors: keiler [email protected]
 * @LastEditTime: 2022-12-07 09:34:32
 * @FilePath: /capture-plugin/include/grab/stream/queue.h
 * @Description: 线程安全的队列
 */
#include <iostream>
#include <thread>
#include <unistd.h>
#include <atomic>
#include <condition_variable>
#include <functional>
#include <memory>
#include <mutex>
#include <queue>

namespace Stream
{
    template <class T>
    class Queue
    {
    protected:
        // Data
        std::queue<T> _queue;
        typename std::queue<T>::size_type _size_max;
        // Thread gubbins
        std::mutex _mutex;
        std::condition_variable _fullQue;
        std::condition_variable _empty;
        // Exit
        // 原子操作
        std::atomic_bool _quit;     //{ false };
        std::atomic_bool _finished; // { false };

    public:
        Queue(const size_t size_max) : _size_max(size_max)
        {
            _quit = ATOMIC_VAR_INIT(false);
            _finished = ATOMIC_VAR_INIT(false);
        }

        bool push(T &data)
        {
            std::unique_lock<std::mutex> lock(_mutex);
            while (!_quit && !_finished)
            {
                if (_queue.size() < _size_max)
                {
                    _queue.push(std::move(data));
                    //_queue.push(data);
                    _empty.notify_all();
                    return true;
                }
                else
                {
                    // wait的时候自动释放锁，如果wait到了会获取锁
                    _fullQue.wait(lock);
                }
            }
            return false;
        }

        bool pop(T &data)
        {
            std::unique_lock<std::mutex> lock(_mutex);
            while (!_quit)
            {
                if (!_queue.empty())
                {
                    // data = std::move(_queue.front());
                    data = _queue.front();
                    _queue.pop();
                    _fullQue.notify_all();
                    return true;
                }
                else if (_queue.empty() && _finished)
                {
                    return false;
                }
                else
                {
                    _empty.wait(lock);
                }
            }
            return false;
        }
        // The queue has finished accepting input
        void finished()
        {
            _finished = true;
            _empty.notify_all();
        }
        void quit()
        {
            _quit = true;
            _empty.notify_all();
            _fullQue.notify_all();
        }
        int length()
        {
            return static_cast<int>(_queue.size());
        }
    };
}