save frames

saveframes.c

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libavutil/imgutils.h>
#include <libavutil/pixfmt.h>
#include <libswscale/swscale.h>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"


void save_frame(AVFrame * frame, int w, int h, int i)
{
    char filename[24];

    sprintf(filename, "frame%04d.png", i);

    stbi_write_png(filename, w, h, 3, frame->data[0], frame->linesize[0]);
}

int main(int argc, char *argv[])
{
    AVFormatContext *fmt_ctx = NULL;
    AVCodecContext *dec_ctx = NULL;
    AVCodec *dec = NULL;
    AVFrame *frame = NULL;
    AVFrame *frame_rgb = NULL;
    struct SwsContext *sws_ctx = NULL;
    int ret;
    int video_stream_index;
    int width, height;
    enum AVPixelFormat pix_fmt;
    uint8_t *video_dst_data[4] = { NULL };
    int video_dst_linesize[4] = { 0 };
    int video_dst_bufsize;
    uint8_t * buffer = NULL;
    int num_bytes;

    av_register_all();

    if ((ret = avformat_open_input(&fmt_ctx, argv[1], NULL, NULL)) < 0) {
        fprintf(stderr, "Failed to open file %s\n", argv[1]);
        return -1;
    }

    if ((ret = avformat_find_stream_info(fmt_ctx, NULL)) < 0) {
        fprintf(stderr, "failed to find stream info %s\n", argv[1]);
        return -1;
    }

    if ((ret = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &dec, 0)) < 0) {
        fprintf(stderr, "No streams were found\n");
        return -1;
    }

    video_stream_index = ret;
    dec_ctx = fmt_ctx->streams[video_stream_index]->codec;

    width = dec_ctx->width;
    height = dec_ctx->height;
    pix_fmt = dec_ctx->pix_fmt;

    if ((ret = av_image_alloc(video_dst_data, video_dst_linesize, width, height, pix_fmt, 1)) < 0) {
        fprintf(stderr, "Could not allocate image data buffer\n");
        return -1;
    }
    video_dst_bufsize = ret;

    av_dump_format(fmt_ctx, 0, argv[1], 0);

    if ((ret = avcodec_open2(dec_ctx, dec, NULL)) < 0) {
        fprintf(stderr, "Could not open codec\n");
        return -1;
    }
 
    if (!(frame = av_frame_alloc())) {
        fprintf(stderr, "Could not alloc frame\n");
        return -1;
    }

    if (!(frame_rgb = av_frame_alloc())) {
        fprintf(stderr, "Could not alloc frame\n");
        return -1;
    }

    num_bytes = avpicture_get_size(AV_PIX_FMT_RGB24, dec_ctx->width, dec_ctx->height);
    buffer = (uint8_t *)av_malloc(num_bytes * sizeof(uint8_t));

    avpicture_fill((AVPicture *) frame_rgb, buffer, AV_PIX_FMT_RGB24, dec_ctx->width, dec_ctx->height);

    sws_ctx = sws_getContext(
            dec_ctx->width,
            dec_ctx->height,
            dec_ctx->pix_fmt,
            dec_ctx->width,
            dec_ctx->height,
            AV_PIX_FMT_RGB24,
            SWS_BILINEAR,
            NULL,
            NULL,
            NULL
    );

    int i = 0;
    while(1) {
        AVPacket packet;

        if ((ret = av_read_frame(fmt_ctx, &packet)) < 0) {
            fprintf(stderr, "Could not read frame\n");
            break;
        }

        if (packet.stream_index == video_stream_index) {
            int got_frame = 0;
            
            if ((ret = avcodec_decode_video2(dec_ctx, frame, &got_frame, &packet)) < 0) {
                fprintf(stderr, "Could not decode frame\n");
                break;
            }

            if (got_frame) {
                frame->pts = av_frame_get_best_effort_timestamp(frame);
               
                // Do something with frame
                //
                int height = sws_scale(sws_ctx, (uint8_t const * const *)frame->data,
                          frame->linesize, 0, frame->height,
                          frame_rgb->data, frame_rgb->linesize);

                // NOW: frame_rgb->data holds the rgb data

                save_frame(frame_rgb, dec_ctx->width, dec_ctx->height, i++);
                //av_frame_unref(frame);
            } 
        } 

        av_packet_unref(&packet);
    }

    return 0;
}