Libav mp4 to mpeg-ts remuxing

/*
 *
 * SAMPLE CODE FOR MP4-to-MPEG-TS CONTAINER CONVERSION
 *
 * This is the basis for a python plugin for doing on-the-fly segmentation of
 *
 *
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include <libavformat/avformat.h>
#include <libavutil/mathematics.h>

typedef struct config_info
{
    const char *input_filename;
    const char *output_playlist;
    const char *output_prefix;
    int segment_length;
} config_info;

static AVStream *add_output_stream(AVFormatContext *output_format_context, AVStream *input_stream)
{
    AVCodecContext *input_codec_context;
    AVCodecContext *output_codec_context;
    AVStream *output_stream;

    output_stream = avformat_new_stream(output_format_context, NULL);
    if (!output_stream) {
        fprintf(stderr, "Could not allocate stream\n");
        exit(1);
    }

    input_codec_context = input_stream->codec;
    output_codec_context = output_stream->codec;

    output_codec_context->codec_id = input_codec_context->codec_id;
    output_codec_context->codec_type = input_codec_context->codec_type;
    output_codec_context->codec_tag = input_codec_context->codec_tag;
    output_codec_context->bit_rate = input_codec_context->bit_rate;
    output_codec_context->extradata = input_codec_context->extradata;
    output_codec_context->extradata_size = input_codec_context->extradata_size;

    // Some codecs put the wrong ticks_per_frame in (in the thousands instead of tens)
    if(av_q2d(input_codec_context->time_base) * input_codec_context->ticks_per_frame > av_q2d(input_stream->time_base) && av_q2d(input_stream->time_base) < 1.0/1000) {
        output_codec_context->time_base = input_codec_context->time_base;
        output_codec_context->time_base.num *= input_codec_context->ticks_per_frame;
    }
    else {
        output_codec_context->time_base = input_stream->time_base;
    }

    switch (input_codec_context->codec_type) {
        case AVMEDIA_TYPE_AUDIO:
            output_codec_context->channel_layout = input_codec_context->channel_layout;
            output_codec_context->sample_rate = input_codec_context->sample_rate;
            output_codec_context->channels = input_codec_context->channels;
            output_codec_context->frame_size = input_codec_context->frame_size;
            if ((input_codec_context->block_align == 1 && input_codec_context->codec_id == AV_CODEC_ID_MP3) || input_codec_context->codec_id == AV_CODEC_ID_AC3) {
                output_codec_context->block_align = 0;
            }
            else {
                output_codec_context->block_align = input_codec_context->block_align;
            }
            fprintf(stdout, "Audio time bases:\n");
            break;
        case AVMEDIA_TYPE_VIDEO:
            output_codec_context->pix_fmt = input_codec_context->pix_fmt;
            output_codec_context->width = input_codec_context->width;
            output_codec_context->height = input_codec_context->height;
            output_codec_context->has_b_frames = input_codec_context->has_b_frames;

            if (output_format_context->oformat->flags & AVFMT_GLOBALHEADER) {
                output_codec_context->flags |= CODEC_FLAG_GLOBAL_HEADER;
            }
            fprintf(stdout, "Video time bases:\n");
            break;
        default:
            break;
    }

    fprintf(stdout, "- Input Codec: %d\n- Output Codec: %d\n", input_codec_context->time_base, output_codec_context->time_base);
    fprintf(stdout, "- Input Stream: %d\n- Output Stream: %d\n", input_stream->time_base, output_stream->time_base);

    return output_stream;
}

int main(int argc, char **argv)
{
    if (argc != 5)
    {
        fprintf(stderr, "Usage: %s <input filename> <playlist filename> <output prefix> <segment length>\n", argv[0]);
        exit(1);
    }

    config_info config;
    memset(&config, 0, sizeof(config_info));

    config.input_filename = argv[1];
    config.output_playlist = argv[2];
    config.output_prefix = argv[3];
    config.segment_length = atoi(argv[4]);

    char *output_filename = (char*) malloc(sizeof(char) * (strlen(config.output_prefix) + 10));
    if (!output_filename)
    {
        fprintf(stderr, "Segmenter error: Could not allocate space for output filenames with prefix %s\n", config.output_prefix);
        exit(1);
    }

    avcodec_register_all();
    av_register_all();

    AVInputFormat *input_format = av_find_input_format("mp4");
    if (!input_format)
    {
        fprintf(stderr, "Segmenter error: Could not find MP4 demuxer.\n");
        exit(1);
    }

    AVFormatContext *input_context = NULL;
    int ret = avformat_open_input(&input_context, config.input_filename, input_format, NULL);
    if (ret != 0)
    {
        fprintf(stderr, "Segmenter error: Could not open input file, make sure it is an MP4 file: %d\n", ret);
        exit(1);
    }

    if (avformat_find_stream_info(input_context, NULL) < 0)
    {
        fprintf(stderr, "Segmenter error Could not read stream information.\n");
        exit(1);
    }

    AVOutputFormat *output_format = av_guess_format("mpegts", NULL, NULL);
    if (!output_format)
    {
        fprintf(stderr, "Segmenter error: Could not find MPEG-TS muxer.\n");
        exit(1);
    }

    AVFormatContext *output_context = avformat_alloc_context();
    if (!output_context)
    {
        fprintf(stderr, "Segmenter error: Could not allocate output context.\n");
        exit(1);
    }
    output_context->oformat = output_format;

    int video_index = -1;
    int audio_index = -1;

    AVStream *video_stream;
    AVStream *audio_stream;

    int i;

    for (i = 0; i < input_context->nb_streams && (video_index < 0 || audio_index < 0); i++)
    {
        switch (input_context->streams[i]->codec->codec_type)
        {
            case AVMEDIA_TYPE_VIDEO:
                video_index = i;
                input_context->streams[i]->discard = AVDISCARD_NONE;
                video_stream = add_output_stream(output_context, input_context->streams[i]);
                break;

            case AVMEDIA_TYPE_AUDIO:
                audio_index = i;
                input_context->streams[i]->discard = AVDISCARD_NONE;
                audio_stream = add_output_stream(output_context, input_context->streams[i]);
                break;

            default:
                input_context->streams[i]->discard = AVDISCARD_ALL;
                break;
        }
    }

    if (video_index == -1 || audio_index == -1)
    {
        fprintf(stderr, "Segmenter error: Video or audio stream not found.\n");
        exit(1);
    }

    av_dump_format(input_context, 0, config.input_filename, 0);
    av_dump_format(output_context, 0, config.output_prefix, 1);

    AVCodec *codec = avcodec_find_decoder(video_stream->codec->codec_id);
    if (!codec)
    {
        fprintf(stderr, "Segmenter error: Could not find video decoder, key frames will not be honored.\n");
    }

    if (avcodec_open2(video_stream->codec, codec, NULL) < 0)
    {
        fprintf(stderr, "Segmenter error: Could not open video decoder, key frames will not be honored.\n");
    }

    AVCodec *codec2 = avcodec_find_decoder(audio_stream->codec->codec_id);
    if (!codec2)
    {
        fprintf(stderr, "Segmenter error: Could not find audio decoder, key frames will not be honored.\n");
    }

    if (avcodec_open2(audio_stream->codec, codec2, NULL) < 0)
    {
        fprintf(stderr, "Segmenter error: Could not open audio decoder, key frames will not be honored.\n");
    }

    // Generate output filename
    snprintf(output_filename, strlen(config.output_prefix) + 11, "%s_result.ts", config.output_prefix);

    if (avio_open(&output_context->pb, output_filename, AVIO_FLAG_WRITE) < 0)
    {
        fprintf(stderr, "Segmenter error: Could not open '%s'\n", output_filename);
        exit(1);
    }

    if (avformat_write_header(output_context, NULL))
    {
        fprintf(stderr, "Segmenter error: Could not write mpegts header to first output file.\n");
        exit(1);
    }

    // Required when muxing h264 into MPEG-TS
    AVBitStreamFilterContext *bsf = av_bitstream_filter_init("h264_mp4toannexb");
    if (!bsf)
    {
        fprintf(stderr, "Error initializing the h264_mp4toannexb filter!");
        exit(1);
    }

    int decode_done;
    do
    {
        AVPacket packet;
        av_init_packet(&packet);
        decode_done = av_read_frame(input_context, &packet);
        if (decode_done < 0)
        {
            break;
        }

        if (av_dup_packet(&packet) < 0)
        {
            fprintf(stderr, "Segmenter error: Could not duplicate packet (av_dup_packet).\n");
            av_free_packet(&packet);
            break;
        }

        if (packet.stream_index == video_index)
        {
            int a = av_bitstream_filter_filter(bsf, video_stream->codec, NULL,
                                               &packet.data, &packet.size,
                                               packet.data, packet.size,
                                               packet.flags & AV_PKT_FLAG_KEY);

            if (a == 0)
            {
                fprintf(stderr, "Eh?!\n");
            }
            else if (a < 0)
            {
                fprintf(stderr, "%s failed for stream %d, codec %s, error %d\n", bsf->filter->name, packet.stream_index, video_stream->codec->codec->name, a);
                exit(1);
            }
        }

        // Correct time base stuff
        // TOOODOOOOO!

        ret = av_interleaved_write_frame(output_context, &packet);
        if (ret < 0)
        {
            fprintf(stderr, "Segmenter error: Could not write frame of stream: %d\n", ret);
            exit(1);
        }
        else if (ret > 0)
        {
            fprintf(stderr, "Segmenter info: End of stream requested.\n");
            av_free_packet(&packet);
            break;
        }

        av_free_packet(&packet);
    } while (!decode_done);

    av_write_trailer(output_context);
    avcodec_close(video_stream->codec);

    for (i = 0; i < output_context->nb_streams; i++)
    {
        av_freep(&output_context->streams[i]->codec);
        av_freep(&output_context->streams[i]);
    }

    av_bitstream_filter_close(bsf);

    avio_close(output_context->pb);
    av_free(output_context);

    fprintf(stdout, "SUCCESS!!!\n");
    exit(0);
}