ffmpeg从mp4文件中提取yuv420p图片

45 阅读2分钟

使用C语言从mp4文件中提取yuv420p格式的图片,方便后实验。

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libgen.h>
#include <sys/stat.h>

 
int main(int argc, char **argv)
{
    AVFormatContext *ifmt_ctx = NULL;
    AVCodecContext *codec_ctx = NULL;
    AVCodecParameters *codec_para = NULL;
    AVPacket *pkt = NULL;
    AVFrame *frame = NULL;

    char *in_filename, *out_filename;
    int ret, i;
    int video_stream_index = -1;
    int cnt = 0;

    char out_file[256];
    in_filename  = "../input.mp4";
    out_filename = "frame_out/%dx%d_%d.yuv420p";

    // 创建目录
    mkdir(dirname(out_filename), 0777);

    printf("out_filename: %s\n", out_filename);
 
    pkt = av_packet_alloc();
    if (!pkt) {
        fprintf(stderr, "Could not allocate AVPacket\n");
        goto end;
    }

    frame = av_frame_alloc();
    if (!frame) {
        fprintf(stderr, "Could not allocate AVFrame\n");
        goto end;
    }

    if ((ret = avformat_open_input(&ifmt_ctx, in_filename, 0, 0)) < 0) {
        fprintf(stderr, "Could not open input file '%s'", in_filename);
        goto end;
    }
 
    if ((ret = avformat_find_stream_info(ifmt_ctx, 0)) < 0) {
        fprintf(stderr, "Failed to retrieve input stream information");
        goto end;
    }
 
    av_dump_format(ifmt_ctx, 0, in_filename, 0);

    for (i = 0; i < ifmt_ctx->nb_streams; i++) {
        if (ifmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_index = i;
            break;
        }
    }
 
    codec_para = ifmt_ctx->streams[video_stream_index]->codecpar;
    const AVCodec *codec = avcodec_find_decoder(codec_para->codec_id);
    if (codec == NULL) {
        fprintf(stderr, "can not find video decoder\n");
        goto end;
    }

    codec_ctx = avcodec_alloc_context3(codec);
    if (codec_ctx == NULL) {
        fprintf(stderr, "can not alloc context\n");
        goto end;
    }
    avcodec_parameters_to_context(codec_ctx, codec_para);

    ret = avcodec_open2(codec_ctx, NULL, NULL);
    if (ret < 0) {
        fprintf(stderr, "can not open context\n");
        goto end;
    }

    int w = codec_ctx->width;
    int h = codec_ctx->height;

    printf("video_stream_index: %d\n", video_stream_index);
    printf("width: %d\n", w);
    printf("height: %d\n", h);

    printf("nb_frames: %lld\n", ifmt_ctx->streams[video_stream_index]->nb_frames);
    printf("avg_frame_rate.num: %d\n", ifmt_ctx->streams[video_stream_index]->avg_frame_rate.num);
    printf("avg_frame_rate.den: %d\n", ifmt_ctx->streams[video_stream_index]->avg_frame_rate.den);

    // 平均帧率
    int avg_frame_rate = ifmt_ctx->streams[video_stream_index]->avg_frame_rate.num / ifmt_ctx->streams[video_stream_index]->avg_frame_rate.den;
    printf("avg_frame_rate: %d\n", avg_frame_rate);

    while (1) {
        AVStream *in_stream;
 
        ret = av_read_frame(ifmt_ctx, pkt);
        if (ret < 0)
            break;
 
        in_stream  = ifmt_ctx->streams[pkt->stream_index];
        if (pkt->stream_index != video_stream_index) {
            av_packet_unref(pkt);
            continue;
        }

        if (avcodec_send_packet(codec_ctx, pkt) == 0) {
          while (avcodec_receive_frame(codec_ctx, frame) == 0) {

            // 由于frame数据量太大,所以只每隔1秒保存一张yuv420p的图片
            if (cnt % avg_frame_rate == 0) {
                snprintf(out_file, sizeof(out_file), out_filename, w, h, cnt);
                FILE * f = fopen(out_file, "wb");

                fwrite(frame->data[0], 1, w*h, f);    // y
                fwrite(frame->data[1], 1, w*h/4, f);  // u
                fwrite(frame->data[2], 1, w*h/4, f);  // v

                fclose(f);
            }

            cnt++;
            av_frame_unref(frame);
          }
        }

        av_packet_unref(pkt);
    }

    // 读取codec_ctx剩余的帧
    if (avcodec_send_packet(codec_ctx, NULL) == 0) {
        while (avcodec_receive_frame(codec_ctx, frame) == 0) {

            // 由于frame数据量太大,所以只每隔1秒保存一张yuv420p的图片
            if (cnt % avg_frame_rate == 0) {
                snprintf(out_file, sizeof(out_file), out_filename, w, h, cnt);
                FILE * f = fopen(out_file, "wb");

                fwrite(frame->data[0], 1, w*h, f);    // y
                fwrite(frame->data[1], 1, w*h/4, f);  // u
                fwrite(frame->data[2], 1, w*h/4, f);  // v

                fclose(f);
            }

            cnt++;
            av_frame_unref(frame);
        }
    }

    printf("Number of frames actually read: %d\n", cnt);

end:
    av_packet_free(&pkt);
    av_frame_free(&frame);
    avcodec_free_context(&codec_ctx);
    avformat_close_input(&ifmt_ctx);

    return 0;
}