imgs2video.c

#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <dirent.h>
#include <sys/stat.h>
#include <assert.h>
#include <errno.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavfilter/avfilter.h>
#include <libavfilter/avfiltergraph.h>
#include <libavfilter/vsrc_buffer.h>
#include <libavutil/avutil.h>
#include <libavutil/imgutils.h>
#ifndef LIBAV
#include <libavfilter/avcodec.h>
#endif

/* resolution must be a multiple of two */
#define WIDTH 20
#define HEIGHT 512
#define N_FRAMES 100
#define OUT_FILE "out.mp4"
#define VCODEC "libx264"
#define FRAME_RATE 25
#define BITRATE (200*1000)
#define PROFILE "main"
#define PRESET "medium"
/* profile, preset are settings of x264 encoder */
#define FILTERS_STRING "scale"
/*
 * "scale" filter gives seamless conversion between pixel formats
 * PIX_FMT_RGBA is the format you want to work with
 * PIX_FMT_YUV420P is the one in which h264 encoder requires
 * (likely all encoders accept YUV format, and not sure if any encoder
 * accepts RGB* formats)
 * This string can contail conveyor of filters, see avfilter docs
 */

struct transcoder {
    AVFormatContext *out;
    AVCodecContext *enc;
    AVFilterContext *filter_src;
    AVFilterContext *filter_sink;
    AVFilterGraph *filter_graph;
    uint8_t *video_outbuf;
    int video_outbuf_size;
    int encoded_size; /* number of bytes with actual encoded data */
    unsigned int frames_in; /* how many have been read */
    unsigned int frames_out;
    unsigned int pts_step; /* timestamp step, in stream time_base units */
};
typedef struct transcoder Transcoder;

static int global_init(void);
static int open_out(Transcoder *tc);
static int open_encoder(Transcoder *tc);
static int setup_filters(Transcoder *tc);
static int tc_process_frame(Transcoder *tc, unsigned int i);
static int tc_flush_encoder(Transcoder *tc);

int main(int argc, char **argv) {
    int r;
    Transcoder *tc;
    unsigned int i;

    r = global_init();
    assert(!r);

    tc = calloc(1, sizeof(*tc));
    assert(tc);

    r = open_out(tc);
    if (r) {
        fprintf(stderr, "Open out file fail\n");
        return r;
    }

    r = open_encoder(tc);
    if (r) {
        fprintf(stderr, "Encoder open fail\n");
        return r;
    }

    r = setup_filters(tc);
    if (r) {
        fprintf(stderr, "Filters setup fail\n");
        return r;
    }

    r = avformat_write_header(tc->out, NULL);
    if (r) {
        fprintf(stderr, "write out file fail\n");
        return r;
    }
    AVRational ratio = av_div_q(tc->out->streams[0]->time_base, tc->enc->time_base);
    tc->pts_step = ratio.den / ratio.num;
    printf("pts_step %d\n", tc->pts_step);

    for(i = 0; i < N_FRAMES; i++) {
        r = tc_process_frame(tc, i);
        if (r < 0) {
            fprintf(stderr, "Fatal error processing frame, aborting\n");
            break;
        }
    }
    tc_flush_encoder(tc);

    av_write_trailer(tc->out);
    av_dump_format(tc->out, 0, OUT_FILE, 1);

    avcodec_close(tc->enc);
    av_free(tc->video_outbuf);
    avio_close(tc->out->pb);
    avformat_free_context(tc->out);
    free(tc);

    return 0;
}

#ifdef LIBAV
/* add missing proc for compatibility */
static int avfilter_fill_frame_from_video_buffer_ref(AVFrame *frame,
        const AVFilterBufferRef *picref)
{
    if (!picref || !picref->video || !frame)
        return AVERROR(EINVAL);

    memcpy(frame->data,     picref->data,     sizeof(frame->data));
    memcpy(frame->linesize, picref->linesize, sizeof(frame->linesize));
    frame->interlaced_frame = picref->video->interlaced;
    frame->top_field_first  = picref->video->top_field_first;
    frame->key_frame        = picref->video->key_frame;
    frame->pict_type        = picref->video->pict_type;

    return 0;
}
#endif

static int global_init(void) {
    av_log_set_level(AV_LOG_DEBUG);
    av_register_all();
    avfilter_register_all();
    return 0;
}

static int open_out(Transcoder *tc) {
    /* auto detect the output format from the name */
    AVOutputFormat *fmt;
    fmt = av_guess_format(NULL, OUT_FILE, NULL);
    if (!fmt) {
        fprintf(stderr, "Could not find suitable output format\n");
        return 1;
    }
    /* allocate the output media context */
    tc->out = avformat_alloc_context();
    if (!tc->out) {
        fprintf(stderr, "Memory error\n");
        return 1;
    }
    tc->out->oformat = fmt;
    snprintf(tc->out->filename, sizeof(tc->out->filename), "%s", OUT_FILE);

    if (avio_open2(&tc->out->pb, tc->out->filename, AVIO_FLAG_WRITE, NULL, NULL) < 0) {
        fprintf(stderr, "Could not open '%s'\n", tc->out->filename);
        return 1;
    }
    return 0;
}

static int open_encoder(Transcoder *tc) {
    AVCodec *codec = avcodec_find_encoder_by_name(VCODEC);
    if (!codec) {
        fprintf(stderr, "Encoder %s not found\n", VCODEC);
        return 1;
    }
    AVStream *st;
    st = avformat_new_stream(tc->out, codec);
    if (!st) {
        fprintf(stderr, "Could not alloc stream\n");
        return 1;
    }

    st->sample_aspect_ratio.den = 1;
    st->sample_aspect_ratio.num = 1;

    tc->enc = st->codec;
    tc->enc->width = WIDTH;
    tc->enc->height = HEIGHT;

    tc->enc->time_base.den = FRAME_RATE;
    tc->enc->time_base.num = 1;
    tc->enc->pix_fmt = PIX_FMT_YUV420P;

    tc->enc->sample_aspect_ratio.den = 1;
    tc->enc->sample_aspect_ratio.num = 1;

    tc->enc->bit_rate = BITRATE;
    tc->enc->bit_rate_tolerance = tc->enc->bit_rate / 5;
    tc->enc->thread_count = 0; /* use several threads for encoding */

    AVDictionary *opts = NULL;
    if (tc->enc->bit_rate != 0) /* profiles don't support lossless */
        av_dict_set(&opts, "profile", PROFILE, 0);
    else
        av_dict_set(&opts, "qp", "0", 0); /* set lossless mode */
    av_dict_set(&opts, "preset", PRESET, 0);
    /* open the codec */
    if (avcodec_open2(tc->enc, codec, &opts) < 0) {
        fprintf(stderr, "could not open codec\n");
        return 1;
    }

    /* allocate output buffer */
    tc->video_outbuf_size = WIDTH * HEIGHT * 4; /* upper bound */
    tc->video_outbuf = av_malloc(tc->video_outbuf_size);
    if (!tc->video_outbuf) {
        fprintf(stderr, "Alloc outbuf fail\n");
        return 1;
    }
    av_dump_format(tc->out, 0, OUT_FILE, 1);
    return 0;
}

static int setup_filters(Transcoder *tc) {
    int r;
    AVFilter *src  = avfilter_get_by_name("buffer");
    assert(src);
    AVFilter *sink = avfilter_get_by_name("nullsink");
    assert(sink);
    AVFilterInOut *outputs = av_mallocz(sizeof(AVFilterInOut));
    assert(outputs);
    AVFilterInOut *inputs  = av_mallocz(sizeof(AVFilterInOut));
    assert(inputs);
    tc->filter_graph = avfilter_graph_alloc();
    assert(tc->filter_graph);

    char filter_args[50];
    /* buffer video source: the decoded frames from the decoder will be inserted here. */
    snprintf(filter_args, sizeof(filter_args), "%d:%d:%d:%d:%d:%d:%d",
            tc->enc->width, tc->enc->height, PIX_FMT_RGBA,
            tc->enc->time_base.num, tc->enc->time_base.den,
            tc->enc->sample_aspect_ratio.num, tc->enc->sample_aspect_ratio.den);
    r = avfilter_graph_create_filter(&tc->filter_src, src, "in",
            filter_args, NULL, tc->filter_graph);
    if (r < 0) {
        av_log(NULL, AV_LOG_ERROR, "Cannot create buffer source\n");
        return r;
    }

    /* buffer video sink: to terminate the filter chain. */
    enum PixelFormat pix_fmts[] = { PIX_FMT_YUV420P, PIX_FMT_NONE };
    r = avfilter_graph_create_filter(&tc->filter_sink, sink, "out",
            NULL, pix_fmts, tc->filter_graph);
    if (r < 0) {
        av_log(NULL, AV_LOG_ERROR, "Cannot create buffer sink\n");
        return r;
    }

    /* Endpoints for the filter graph. */
    outputs->name       = av_strdup("in");
    outputs->filter_ctx = tc->filter_src;

    inputs->name       = av_strdup("out");
    inputs->filter_ctx = tc->filter_sink;

    char *filter_descr = FILTERS_STRING;

    r = avfilter_graph_parse(tc->filter_graph, filter_descr,
#ifdef LIBAV
            inputs, outputs,
#else
            &inputs, &outputs,
#endif
            NULL);

    if (r < 0) {
        printf("avfilter_graph_parse fail\n");
        return r;
    }

    if ((r = avfilter_graph_config(tc->filter_graph, NULL)) < 0) {
        printf("avfilter_graph_config fail\n");
        return r;
    }

    return 0;
}

static int tc_process_frame_input(Transcoder *tc, unsigned int i);
static int tc_process_frame_output(Transcoder *tc);

/**
 * @return 0 on success, <0 on fatal error, 1 on non-fatal error
 */
static int tc_process_frame(Transcoder *tc, unsigned int i) {
    int r;
    r = tc_process_frame_input(tc, i);
    if (r) {
        fprintf(stderr, "tc_process_frame_input fail\n");
        return r;
    }
    r = tc_process_frame_output(tc);
    if (r) {
        fprintf(stderr, "tc_process_frame_output fail\n");
        return r;
    }
    return 0;
}

static int tc_write_encoded(Transcoder *tc);

static int tc_flush_encoder(Transcoder *tc) {
    int r;
    while (1) {
        /* flush buffered remainings */
        r = avcodec_encode_video(tc->enc, tc->video_outbuf, tc->video_outbuf_size, NULL);
        if (r <= 0)
            break;
        tc->encoded_size = r;
        tc_write_encoded(tc);
    }
    return 0;
}

static int draw_frame(Transcoder *tc, AVFrame *f) {
    int r;
    avcodec_get_frame_defaults(f);
    r = avpicture_alloc((AVPicture*)f, PIX_FMT_RGBA, tc->enc->width, tc->enc->height);
    /* All the data is accessible by f->data[0] pointer
     * f->linesize[0] shows number of bytes used to represent single row of pixels */
    f->width = tc->enc->width;
    f->height = tc->enc->height;
    f->format = PIX_FMT_RGBA;
    f->pts = tc->frames_in++;
    f->pict_type = 0; /* let codec choose */

    memset(f->data[0], 0, f->width * f->height * 4); /* black */
    int i; /* vertical, y, height */
    int j; /* horizontal, x, width */
    uint8_t *pix_ptr = f->data[0];
    for (i = 0; i < f->height; i++) {
        for (j = 0; j < f->width; j++) {
            pix_ptr[0] = 0xff; /* red is first byte on little-endian. See libavutil/pixfmt.h */
            pix_ptr[1] = i % 0x100;
            pix_ptr[2] = (tc->frames_in * 2) % 0x100;
            pix_ptr += 4;
        }
    }

    return r;
}

/**
 * @return 0 on success, <0 on fatal error, 1 on non-fatal error
 */
static int tc_process_frame_input(Transcoder *tc, unsigned int i) {
    AVFrame *pFrame;
    int r;

    pFrame = avcodec_alloc_frame();
    if (!pFrame) {
        printf("Can't allocate memory for AVFrame\n");
        return 1;
    }

    r = draw_frame(tc, pFrame);
    if (r)
        return r;

    /* push the decoded frame into the filtergraph */
#ifdef LIBAV
    r = av_vsrc_buffer_add_frame(tc->filter_src, pFrame, pFrame->pts, (AVRational){1, 1});
#else
    r = av_vsrc_buffer_add_frame(tc->filter_src, pFrame, 0);
#endif
    assert(r >= 0);
    av_free(pFrame);
    return 0;
}

/**
 * @return 0 on success, <0 on fatal error, 1 on non-fatal error
 */
static int tc_process_frame_output(Transcoder *tc) {
    int r;
    int out_size;

    /* pull filtered pictures from the filtergraph */
    AVFilterBufferRef *picref = NULL;
    r = avfilter_poll_frame(tc->filter_sink->inputs[0]);
    if (r < 0) {
        fprintf(stderr, "avfilter_poll_frame fail %d\n", r);
        return -1;
    }
    if (r == 0) {
        printf("avfilter_poll_frame: no frames available\n");
        return 0;
    }

    r = avfilter_request_frame(tc->filter_sink->inputs[0]);
    if (r) {
        fprintf(stderr, "avfilter_request_frame fail %d\n", r);
        return -1;
    }

    picref = tc->filter_sink->inputs[0]->cur_buf;
    assert(picref);

    AVFrame *picture = avcodec_alloc_frame();
    assert(picture);
    r = avfilter_fill_frame_from_video_buffer_ref(picture, picref);
    assert(r == 0);
    picture->pts = picref->pts;
    /* encode the image */
    out_size = avcodec_encode_video(tc->enc, tc->video_outbuf, tc->video_outbuf_size, picture);
    avfilter_unref_buffer(picref);

    /* if zero size, it means the image was buffered */
    if (out_size == 0) {
        printf("encoded frame, no data out, filling encoder buffer\n");
        return 0;
    }
    tc->encoded_size = out_size;
    r = tc_write_encoded(tc);
    if (r) {
        fprintf(stderr, "Error while writing video frame\n");
        return -1;
    }
    return 0;
}

static int tc_write_encoded(Transcoder *tc) {
    int ret;
    /* write from internal buffer */
    AVPacket pkt;
    av_init_packet(&pkt);

    pkt.pts = tc->frames_out++ * tc->pts_step;
    pkt.dts = pkt.pts;
    pkt.duration = tc->pts_step;
    if(tc->enc->coded_frame->key_frame)
        pkt.flags |= AV_PKT_FLAG_KEY;
    pkt.data = tc->video_outbuf;
    pkt.size = tc->encoded_size;

    /* write the compressed frame in the media file */
    ret = av_interleaved_write_frame(tc->out, &pkt);
    av_free_packet(&pkt);
    return ret;
}