使用ffmpeg sdk保存rtmp流为jpeg

1,119 阅读1分钟

主要函数为:

//函数作用:将解码出来的YUV数据保存成JPG图像
int EncodeAndSaveJPEG(const char* out_file_name, AVFrame* pYUVFrame)
{
    AVFormatContext* pFormatCtx;
    AVCodecContext* pCodecCtx;
    AVCodec* pCodec;
    AVStream* video_st;
    AVPacket pkt;
    int ret = 0;

    //根据文件后缀猜一个Format
    avformat_alloc_output_context2(&pFormatCtx, NULL, NULL, out_file_name);

    pCodecCtx = avcodec_alloc_context3(NULL);
    pCodecCtx->codec_id = pFormatCtx->oformat->video_codec;//编码ID
    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    pCodecCtx->pix_fmt = AV_PIX_FMT_YUVJ420P;
    pCodecCtx->width = pYUVFrame->width;
    pCodecCtx->height = pYUVFrame->height;
    pCodecCtx->time_base.num = 1;
    pCodecCtx->time_base.den = 25;

    pCodec = avcodec_find_encoder(pCodecCtx->codec_id);
    if (!pCodec) {
        fprintf(stderr, "jpeg Codec not found. \n");
        return -1;
    }
    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) {
        fprintf(stderr, "Could not open jpeg codec. \n");
        return -1;
    }

    //设置JPEG质量
    pCodecCtx->qcompress = 1; // 0~1.0, default is 0.5
    pCodecCtx->qmin = 2;
    pCodecCtx->qmax = 31;
    pCodecCtx->max_qdiff = 3;

    //给媒体文件创建一股流
    //avformat_new_stream() should be called by the user before avformat_write_header()
    video_st = avformat_new_stream(pFormatCtx, NULL);
    if (video_st == NULL) {
        return -1;
    }

    //Output some information
    av_dump_format(pFormatCtx, 0, out_file_name, 1);

    //Write Header
    avformat_write_header(pFormatCtx, NULL);

    //预申请mjpeg压缩后的packet大小
    av_new_packet(&pkt, pCodecCtx->width * pCodecCtx->height * 3);

    //Encode
    ret = avcodec_send_frame(pCodecCtx, pYUVFrame);
    if (ret == 0)
    {
        while (avcodec_receive_packet(pCodecCtx, &pkt) == 0) {
            ret = av_write_frame(pFormatCtx, &pkt);
        }
    }

    av_packet_unref(&pkt);

    //Write Trailer
    av_write_trailer(pFormatCtx);

    avio_close(pFormatCtx->pb);
    avformat_free_context(pFormatCtx);
    avcodec_close(pCodecCtx);
    avcodec_free_context(&pCodecCtx);

    return 0;
}

主函数就是读取rtmp流,并通过x264解码为YUV avframe

#include <iostream>

#ifdef __cplusplus
extern "C"
{
#endif
#include "libavformat/avformat.h"
#include "libavutil/mathematics.h"
#include "libavutil/imgutils.h"
#include "libavutil/time.h"
#ifdef __cplusplus
}
#endif

#pragma comment(lib, "ffmpeg/libswscale.a")
#pragma comment(lib, "ffmpeg/libavutil.a")
#pragma comment(lib, "ffmpeg/libavcodec.a")
#pragma comment(lib, "ffmpeg/libavformat.a")
#pragma comment(lib, "ffmpeg/libswresample.a")
#pragma comment(lib, "Secur32.lib")//_AcquireCredentialsHandleA@36 _EncryptMessage@16
#pragma comment(lib,"ws2_32.lib")
#pragma comment(lib, "bcrypt.lib")//_BCryptOpenAlgorithmProvider@16
#pragma comment(lib, "Mfplat.lib")//_MFCreateMediaType@4
#pragma comment(lib, "Mfuuid.lib")//_IID_ICodecAPI _IID_IMFMediaEventGenerator
#pragma comment(lib, "Strmiids.lib")//_IID_ICodecAPI

//////fdk-aac
#pragma comment(lib,"fdkaac/fdk-aac.lib")
///x264
#pragma comment(lib,"x264/libx264.lib")

using namespace std;

int EncodeAndSaveJPEG(const char* out_file_name, AVFrame *pYUVFrame);

int main(int argc, char* argv[])
{
    //Input AVFormatContext
    AVFormatContext* ifmt_ctx = NULL;
    AVDictionary* pOptions = NULL;
    AVPacket pkt;
    AVCodecContext* pCodecCtx;
    AVCodec* pCodec;
    AVFrame* avFrame = av_frame_alloc();
    const char* in_filename;
    int ret;
    unsigned int i;
    int videoindex = -1;
    int frame_index = 0;
    in_filename = "rtmp://192.168.100.58:1935/hls/Sword";

    //Network
    avformat_network_init();

    //Input,打开网络流或文件流
    if ((ret = avformat_open_input(&ifmt_ctx, in_filename, 0, 0)) < 0) {
        printf("Could not open input file.");
        goto end;
    }

    //查找码流信息
    //设置查找时间以避免耗时过长
    //https://blog.csdn.net/a1317338022/article/details/78226834
    ifmt_ctx->probesize = 1 * 1024;
    ifmt_ctx->max_analyze_duration = 1 * AV_TIME_BASE;
    if ((ret = avformat_find_stream_info(ifmt_ctx, &pOptions)) < 0) {
        printf("Failed to retrieve input stream information");
        goto end;
    }

    //查找码流中是否有视频流
    for (i = 0; i < ifmt_ctx->nb_streams; i++)
        if (ifmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            videoindex = i;
            break;
        }

    av_dump_format(ifmt_ctx, 0, in_filename, 0);

    pCodecCtx = avcodec_alloc_context3(NULL);
    if (pCodecCtx == NULL)
    {
        printf("Could not allocate AVCodecContext\n");
        goto end;
    }

    ret = avcodec_parameters_to_context(pCodecCtx, ifmt_ctx->streams[videoindex]->codecpar);
    pCodec = avcodec_find_decoder(pCodecCtx->codec_id);	//指向AVCodec的指针.查找解码器
    if (pCodec == NULL)
    {
        printf("Codec not found.\n");
        goto end;
    }
    //打开解码器
    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0)
    {
        printf("Could not open codec.\n");
        goto end;
    }

    //拉流
    while (1)
    {
        //Get an AVPacket
        ret = av_read_frame(ifmt_ctx, &pkt);
        if (ret < 0)
            break;

        //Print to Screen
        if (pkt.stream_index == videoindex) {
            printf("Receive %8d video frames from input URL\n", frame_index);
            frame_index++;

            //h264 -> YUV
            ret = avcodec_send_packet(pCodecCtx, &pkt);
            av_packet_unref(&pkt);
            if (ret == 0) {
                ret = avcodec_receive_frame(pCodecCtx, avFrame);
                if (ret == 0)
                {
                    char szPath[256] = {0};
                    sprintf_s(szPath, _TRUNCATE, "d:\\%d.jpg", frame_index);
                    EncodeAndSaveJPEG(szPath, avFrame);
                }
            }
        }

        av_packet_unref(&pkt);//av_read_frame
    }

end:
    avformat_close_input(&ifmt_ctx);
    avcodec_free_context(&pCodecCtx);

    return 0;
}